from __future__ import annotations from typing import Dict, Mapping, List, Set from enum import Enum class UnificationError(Exception): def __init__(self, msg): super().__init__(msg) class Type: def find(self): return self def unify(self, _): raise NotImplementedError() def subst(self, _): raise NotImplementedError() def check(self): pass class TVarType(Enum): UNDETERMINED = 1 SEQUENCE = 2 RECORD = 5 TUPLE = 6 LIST = 8 def __le__(self, other): if self.__class__ is other.__class__: return (other.value % self.value) == 0 return NotImplemented def unifier(self, other): if self.__class__ is not other.__class__: raise NotImplementedError() if self <= other: return other elif other <= self: return self else: raise UnificationError(f'cannot unify {self} and {other}') class TVar(Type): next_id = 0 def __init__(self, vrange=None): self.type = TVarType.UNDETERMINED self.rank = 0 self.parent = self self.checked = False self.fields = {} self.range = vrange self.id = TVar.next_id TVar.next_id += 1 def check(self): if self.checked: return self.checked = True if self.range is not None: ty = self.find() # maybe we should replace this with explicit eq if ty is not self and ty not in self.range: raise UnificationError( f'{self.id} cannot be substituted by {ty}') def subst(self, mapping: Mapping[int, Type]): # user cannot specify fields... # so this is safe if self.id in mapping: return mapping[self.id] return self def __str__(self): s = self.find() if isinstance(s, TVar): if len(s.fields) > 0: fields = '{' + ', '.join([f'{k}: {v}' for k, v in s.fields.items()]) + '}' else: fields = '' return str(s.id) + fields else: return str(s) def find(self): root = self parent = self.parent while root is not parent and isinstance(parent, TVar): _, parent = root, root.parent = parent, parent.parent if isinstance(parent, TCall): parent = parent.find() return parent def unify(self, other): x = other.find() y = self.find() if x is y: return if isinstance(y, TVar) and isinstance(x, TVar): # unify field type x.type = x.type.unifier(y.type) # unify fields for k, v in y.fields.items(): if k in x.fields: x.fields[k].unify(v) else: x.fields[k] = v # standard union find if x.rank < y.rank: x, y = y, x y.parent = x if x.rank == y.rank: x.rank += 1 elif isinstance(y, TVar): # check fields if isinstance(x, TVirtual): if y.type not in [TVarType.UNDETERMINED, TVarType.RECORD]: raise UnificationError(f'Cannot unify {y} with {x}') for k, v in y.fields.items(): if k not in x.obj.fields: raise UnificationError( f'Cannot unify {y} with {x}') u = x.obj.fields[k] v.unify(u) elif isinstance(x, TObj): if y.type not in [TVarType.UNDETERMINED, TVarType.RECORD]: raise UnificationError(f'Cannot unify {y} with {x}') for k, v in y.fields.items(): if k not in x.fields: raise UnificationError( f'Cannot unify {y} with {x}') u = x.fields[k] v.unify(u) elif isinstance(x, TList): if y.type not in [TVarType.UNDETERMINED, TVarType.SEQUENCE, TVarType.LIST]: raise UnificationError(f'Cannot unify {y} with {x}') for k, v in y.fields.items(): assert isinstance(k, int) v.unify(x.param) elif isinstance(x, TTuple): if y.type not in [TVarType.UNDETERMINED, TVarType.SEQUENCE, TVarType.TUPLE]: raise UnificationError(f'Cannot unify {y} with {x}') for k, v in y.fields.items(): assert isinstance(k, int) if k >= len(x.params): raise UnificationError(f'Cannot unify {y} with {x}') v.unify(x.params[k]) y.parent = x else: y.unify(x) def __eq__(self, other): s = self.find() o = other.find() if not isinstance(s, TVar): return s == o if isinstance(o, TVar): return s.id == o.id return False class FuncArg: def __init__(self, name, typ, is_optional): self.name = name self.typ = typ self.is_optional = is_optional def __str__(self): return f'{self.name}: {self.typ}' + ('?' if self.is_optional else '') class TCall(Type): def __init__(self, posargs: List[Type], kwargs: Dict[str, Type], ret: Type): self.calls = [[posargs, kwargs, ret, None]] self.parent = self self.rank = 0 self.checked = False def check(self): if self.checked: return self.checked = True self.calls[0][3].check() def find(self): root = self parent = self.parent while root is not parent and isinstance(parent, TCall): _, parent = root, root.parent = parent, parent.parent if parent.calls[0][3] is None: return parent return parent.calls[0][3] def unify(self, other): y = self.find() if y is not self: y.unify(other) return x = other.find() if x is y: return if isinstance(x, TCall): # standard union find if x.rank < y.rank: x, y = y, x y.parent = x if x.rank == y.rank: x.rank += 1 x.calls += y.calls elif isinstance(x, TFunc): fn = x for i in range(len(y.calls)): if not x.instantiated: fn = x.instantiate() posargs, kwargs, ret, _ = y.calls[i] c = y.calls[i] c[3] = fn all_args = set(arg.name for arg in fn.args) required = set(arg.name for arg in fn.args if not arg.is_optional) fn.ret.unify(ret) for i, v in enumerate(posargs): arg = fn.args[i] arg.typ.unify(v) if arg.name in required: required.remove(arg.name) for k, v in kwargs.items(): arg = next((arg for arg in fn.args if arg.name == k), None) if arg is None: raise UnificationError(f'Unknown kwarg {k}') if k not in all_args: raise UnificationError(f'Duplicated kwarg {k}') arg.typ.unify(v) if k in required: required.remove(k) all_args.remove(k) if len(required) > 0: raise UnificationError(f'Missing arguments') elif isinstance(other, TVar): other.unify(self) else: raise UnificationError(f'Cannot unify a call with {other}') class TFunc(Type): def __init__(self, args: List[FuncArg], ret: Type, vars: List[TVar]): self.args = args self.ret = ret self.vars = vars self.instantiated = False self.checked = False def check(self): if self.checked: return self.checked = True for arg in self.args: arg.typ.check() self.ret.check() def subst(self, mapping: Mapping[int, Type]): if len(mapping) == 0: return self return TFunc( [FuncArg(arg.name, arg.typ.subst(mapping), arg.is_optional) for arg in self.args], self.ret.subst(mapping), self.vars) def instantiate(self): mapping = {v.id: TVar(v.range) if isinstance(v, TVar) else TVar() for v in self.vars} result = self.subst(mapping) result.instantiated = True return result def __str__(self): return f'({", ".join(str(arg) for arg in self.args)}) -> {self.ret}' def unify(self, other): other = other.find() if other is self: return if isinstance(other, (TVar, TCall)): other.unify(self) elif isinstance(other, TFunc): if len(self.args) != len(other.args): raise UnificationError( f'cannot unify functions with different parameters') self.ret.unify(other.ret) for a, b in zip(self.args, other.args): if a.name != b.name or a.is_optional != b.is_optional: raise UnificationError( f'cannot unify functions with different parameters') a.typ.unify(b.typ) else: raise UnificationError(f'Cannot unify a function with {other}') class TObj(Type): def __init__(self, name: str, fields: Dict[str, Type], params: List[Type], parents=None): self.name = name self.fields = fields self.params = params self.checked = False if parents is None: self.parents = [] else: self.parents = parents def check(self): if self.checked: return self.checked = True for arg in self.fields.values(): arg.check() def subst(self, mapping: Mapping[int, Type]): if len(mapping) == 0: return self new_params = [] changed = False for v in self.params: if isinstance(v, TVar) and v.id in mapping: new_params.append(mapping[v.id]) changed = True else: new_params.append(v) if not changed: return self return TObj(self.name, {k: v.subst(mapping) for k, v in self.fields.items()}, new_params) def unify(self, other): other = other.find() if other is self: return if isinstance(other, TObj): if self.name != other.name: raise UnificationError(f'Cannot unify {self} with {other}') for k in self.fields: self.fields[k].unify(other.fields[k]) elif isinstance(other, TVar): other.unify(self) else: raise UnificationError(f'Cannot unify an object with {other}') def __str__(self): if len(self.params) > 0: p = '[' + ', '.join(str(p) for p in self.params) + ']' else: p = '' return self.name + p def __eq__(self, other): o = other.find() if isinstance(o, TObj): if self.name != o.name: return False for a, b in zip(self.params, o.params): if a != b: return False return True return False class TVirtual(Type): def __init__(self, obj: TObj): self.obj = obj self.checked = False def check(self): if self.checked: return self.checked = True self.obj.check() def __eq__(self, other): o = other.find() if isinstance(o, TVirtual): return self == o return False def unify(self, other): o = other.find() if isinstance(o, TVirtual): self.obj.unify(o.obj) else: raise UnificationError(f'Cannot unify {self} with {o}') def __str__(self): return f'virtual[{self.obj}]' class TList(Type): def __init__(self, param: Type): self.param = param self.checked = False def check(self): if self.checked: return self.checked = True self.param.check() def unify(self, other): other = other.find() if isinstance(other, TVar): other.unify(self) elif isinstance(other, TList): self.param.unify(other.param) else: raise UnificationError(f'Cannot unify list with {other}') def __str__(self): return f'List[{self.param}]' def __eq__(self, other): o = other.find() if isinstance(o, TList): return self.param == o.param return False class TTuple(Type): def __init__(self, params: List[Type]): self.params = params self.checked = False def check(self): if self.checked: return self.checked = True for p in self.params: p.check() def unify(self, other): other = other.find() if isinstance(other, TVar): other.unify(self) elif isinstance(other, TTuple): if len(self.params) != len(other.params): raise UnificationError( 'cannot unify tuples with different length') for a, b in zip(self.params, other.params): a.unify(b) else: raise UnificationError(f'Cannot unify {self} with {other}') def __str__(self): return f'Tuple[{", ".join(str(p) for p in self.params)}]' def __eq__(self, other): o = other.find() if isinstance(o, TTuple): if len(self.params) != len(o.params): return False for a, b in zip(self.params, o.params): if a != b: return False return True return False