From e0051f129bf202f4d5fee6ab443c0d3be35e5a5d Mon Sep 17 00:00:00 2001
From: Fardale <fardale+git@crans.org>
Date: Mon, 7 Oct 2024 19:44:23 +0200
Subject: [PATCH] feat(Scheme): add possiblity to add flags to variable.

---
 lib/common/Scheme.ml | 44 ++++++++++++++++++++++++++++++++++++++------
 1 file changed, 38 insertions(+), 6 deletions(-)

diff --git a/lib/common/Scheme.ml b/lib/common/Scheme.ml
index 4f486ed..c3e96a6 100644
--- a/lib/common/Scheme.ml
+++ b/lib/common/Scheme.ml
@@ -1,8 +1,29 @@
 type t = {
-  vars : Variable.t List.t ;
+  vars : (Variable.t * Type.t) List.t ;
   ty : Type.t ;
 }
 
+type flags = NonArrow | NonTuple
+
+type var_type = Frozen | Flags of Variable.Flags.t
+
+let parse_var_type =
+  let open CCParse in
+  (string "^" >|= fun _ -> Frozen) <|>
+  (many1 ((string ">" >|= fun _ -> NonArrow)
+  <|> (string "*" >|= fun _ -> NonTuple))
+  >|= fun l -> Flags (List.fold_left (fun flags m ->
+                        match m with
+                        | NonArrow -> Variable.Flags.(set non_arrow flags)
+                        | NonTuple -> Variable.Flags.(set non_tuple flags))
+                        Variable.Flags.empty l))
+
+let parse_var =
+  let open CCParse in
+  let* var_type = parse_var_type in
+  let+ name = U.word in
+  (name, var_type)
+
 let of_string env str =
   let pos = String.find ~sub:". " str in
   let bdgs, ty =
@@ -13,30 +34,41 @@ let of_string env str =
   let vars =
     if pos = -1 then
       String.HMap.values_list bdgs
+      |> CCList.map (fun var -> (var, Type.frozen_var env var))
     else
       let str = String.trim @@ String.take pos str in
       let vars =
-        try CCParse.(parse_string_exn @@ sep ~by:space U.word) str
+        try CCParse.(parse_string_exn @@ (
+          let* l = sep ~by:space parse_var in
+          let+ () = eoi in
+          l)) str
         with CCParse.ParseError _ -> invalid_arg "Schema.of_string"
       in
       let vars =
-        try CCList.map (String.HMap.find bdgs) vars
+        try
+          CCList.map (fun (name, var_type) ->
+            let var = String.HMap.find bdgs name in
+            match var_type with
+              | Frozen -> (var, Type.frozen_var env var)
+              | Flags flags ->
+                let fresh_var = Variable.Gen.gen flags env.var_gen in
+                (var, Type.var env fresh_var)
+          ) vars
         with Not_found -> invalid_arg "Schema.of_string"
       in
       vars
   in
-  let vars = CCList.sort_uniq ~cmp:Variable.compare vars in
+  let vars = CCList.sort_uniq ~cmp:(fun (v1, _) (v2, _) -> Variable.compare v1 v2) vars in
   { vars ; ty }
 
 let to_type env t =
   let subst =
     t.vars
-    |> CCList.map (fun var -> var, Type.frozen_var env var)
     |> Variable.Map.of_list
   in
   Subst.apply env subst t.ty
 
 let pp ppf t =
   Fmt.pf ppf "@[<2>%a.@ %a@]"
-    Fmt.(list ~sep:sp Variable.pp) t.vars
+    Fmt.(list ~sep:sp (fun fmt (v, _) -> Variable.pp fmt v)) t.vars
     Type.pp t.ty