small fixes

lib/typeSubstitution.ml

@@ -12,11 +12,12 @@ let singleton id ty = IdentifierMap.singleton id ty
 
 (** Apply substitution to a type *)
 let rec apply subst = function
+  | Type.Int -> Type.Int
   | Type.Var id as t ->
+    (* Look for a substitution in the map *)
     (match IdentifierMap.find_opt id subst with
      | Some ty' -> apply subst ty'
      | None -> t)
-  | Type.Int -> Type.Int
   | Type.Product (ty1, ty2) -> Type.Product (apply subst ty1, apply subst ty2)
   | Type.Arrow (ty1, ty2) -> Type.Arrow (apply subst ty1, apply subst ty2)
 ;;
@@ -27,10 +28,10 @@ let compose s2 s1 =
     (fun _ ty1 ty2 ->
       match ty1, ty2 with
       (* If we have 2, we pick one of them *)
-      | Some ty1, Some _ -> Some (apply s2 ty1)
+      | Some ty1', Some _ -> Some (apply s2 ty1')
       (* If we have 1, we pick the one we have *)
-      | Some ty1, None -> Some (apply s2 ty1)
-      | None, Some ty2 -> Some (apply s2 ty2)
+      | Some ty1', None -> Some (apply s2 ty1')
+      | None, Some ty2' -> Some (apply s2 ty2')
       (* If we have 0, we return nothing *)
       | None, None -> None)
     s1

lib/typeSubstitution.mli

@@ -5,4 +5,4 @@ val apply : t -> Type.t -> Type.t
 (* compose s2 s1 : first s1, then s2 *)
 val compose : t -> t -> t
 val empty : t
-val singleton : Identifier.t -> Type.t -> Type.t Map.Make(Identifier).t
+val singleton : Identifier.t -> Type.t -> t

lib/unification.ml

@@ -1,16 +1,19 @@
 (** Unify 2 types and, if exists, returns the substitution *)
 let rec unify ty1 ty2 =
   match ty1, ty2 with
-  | Type.Product (p1_ty1, p1_ty2), Type.Product (p2_ty1, p2_ty2)
-  | Type.Arrow (p1_ty1, p1_ty2), Type.Arrow (p2_ty1, p2_ty2) ->
-    (match unify p1_ty1 p2_ty1 with
+  (* Same types *)
+  | Type.Int, Type.Int -> Some TypeSubstitution.empty
+  | Type.Var id1, Type.Var id2 when id1 = id2 -> Some TypeSubstitution.empty
+  (* Different types *)
+  | Type.Var id, ty | ty, Type.Var id -> Some (TypeSubstitution.singleton id ty)
+  | Type.Product (l1, r1), Type.Product (l2, r2) | Type.Arrow (l1, r1), Type.Arrow (l2, r2)
+    ->
+    (match unify l1 l2 with
      | Some s1 ->
-       (match
-          unify (TypeSubstitution.apply s1 p1_ty2) (TypeSubstitution.apply s1 p2_ty2)
-        with
+       (* Apply the substitution *)
+       (match unify (TypeSubstitution.apply s1 r1) (TypeSubstitution.apply s1 r2) with
         | Some s2 -> Some (TypeSubstitution.compose s2 s1)
         | None -> None)
      | None -> None)
-  | ty1, ty2 when ty1 = ty2 -> Some TypeSubstitution.empty
-  | _ -> None
+  | _, _ -> None
 ;;