From bcf6dc1951dda2cf411c0cf2fb5c5e9f9e9a73f5 Mon Sep 17 00:00:00 2001
From: crumbtoo <crumb@disroot.org>
Date: Fri, 5 Apr 2024 15:21:49 -0600
Subject: [PATCH] case expression inference

---
 README.org                                 | 18 +++++-----
 src/Rlp/HindleyMilner.hs                   | 39 +++++++++++++++++++---
 visualisers/hmvis/src/hmvis/annotated.cljs | 12 +++++++
 3 files changed, 56 insertions(+), 13 deletions(-)

diff --git a/README.org b/README.org
index 84bdfd4..f77bfdc 100644
--- a/README.org
+++ b/README.org
@@ -62,15 +62,15 @@ Available debug flags include:
 
 ** DONE [#A] HM memoisation prevents shadowing                             :bug:
    CLOSED: [2024-04-04 Thu 12:29]
-Example:
-#+begin_src haskell
--- >>> runHM' $ infer1 [rlpExpr|let f = \x -> x in f (let f = 2 in f)|]
--- Left [TyErrCouldNotUnify
---         (ConT "Int#")
---         (AppT (AppT FunT (ConT "Int#")) (VarT "$a2"))]
--- >>> :t let f = \x -> x in f (let f = 2 in f)
--- let f = \x -> x in f (let f = 2 in f) :: Int
-#+end_src
+   Example:
+   #+begin_src haskell
+   -- >>> runHM' $ infer1 [rlpExpr|let f = \x -> x in f (let f = 2 in f)|]
+   -- Left [TyErrCouldNotUnify
+   --         (ConT "Int#")
+   --         (AppT (AppT FunT (ConT "Int#")) (VarT "$a2"))]
+   -- >>> :t let f = \x -> x in f (let f = 2 in f)
+   -- let f = \x -> x in f (let f = 2 in f) :: Int
+   #+end_src
 For the time being, I just disabled the memoisation. This is very, very bad.
 *** Closing Remarks
     Fixed by entirely rewriting the type inference algorithm :P. Memoisation is
diff --git a/src/Rlp/HindleyMilner.hs b/src/Rlp/HindleyMilner.hs
index 3521e8a..ec508b4 100644
--- a/src/Rlp/HindleyMilner.hs
+++ b/src/Rlp/HindleyMilner.hs
@@ -15,6 +15,7 @@ import Control.Monad.Accum
 import Control.Monad.Reader
 import Control.Monad
 import Control.Monad.Extra
+import Control.Monad.Free
 import Control.Arrow            ((>>>))
 import Control.Monad.Writer.Strict
 
@@ -40,7 +41,7 @@ import Debug.Trace
 import Data.Functor             hiding (unzip)
 import Data.Functor.Extend
 import Data.Functor.Foldable    hiding (fold)
-import Data.Fix                 hiding (cata, para, cataM)
+import Data.Fix                 hiding (cata, para, cataM, ana)
 import Control.Comonad.Cofree
 import Control.Comonad
 
@@ -136,11 +137,41 @@ gather (InR (LetEF Rec (withoutPatterns -> bs) (te,je))) = do
     elimRecBind (x,(tx,_)) j = elim x tx j
     elimBind (x,(tx,_)) j = elimGenerally x tx j
 
-gather (InR (CaseEF (te,je) [Alter (ConP' n []) (ta,ja)])) = do
-    tc <- freshTv
-    let j = equal te tc <> je <> assume n tc <> ja
+gather (InR (CaseEF (te,je) as)) = do
+    as' <- gatherAlter te `traverse` as
+    t <- freshTv
+    let eqs = allEqual (t : (as' ^.. each . _1))
+        j = je <> foldOf (each . _2) as' <> eqs
+    pure (t,j)
+
+-- gather (InR (CaseEF (te,je) [Alter (ConP' n bs) (ta,ja)])) = do
+--     -- let tc' be the type of the saturated type constructor
+--     tc' <- freshTv
+--     bs <- for bs (\b -> (b ^. singular _VarP,) <$> freshTv)
+--     let tbs = bs ^.. each . _2
+--         tc = foldr (:->) tc' tbs
+--         j = equal te tc' <> je <> assume n tc <> forBinds elim bs ja
+--     pure (ta,j)
+
+gatherAlter :: (Unique :> es)
+            => Type'
+            -> Alter PsName (Type', Judgement)
+            -> Eff es (Type', Judgement)
+gatherAlter te (Alter (ConP' n bs) (ta,ja)) = do
+    -- let tc' be the type of the saturated type constructor
+    tc' <- freshTv
+    bs' <- for bs (\b -> (b ^. singular _VarP,) <$> freshTv)
+    let tbs = bs' ^.. each . _2
+        tc = foldr (:->) tc' tbs
+        j = equal te tc' <> assume n tc <> forBinds elim bs' ja
     pure (ta,j)
 
+allEqual :: [Type'] -> Judgement
+allEqual = fold . ana @[_] \case
+    [] -> Nil
+    [a] -> Nil
+    (a:b:xs) -> Cons (equal a b) (b:xs)
+
 forBinds :: (PsName -> Type' -> Judgement -> Judgement)
          -> [(PsName, Type')] -> Judgement -> Judgement
 forBinds f bs j = foldr (uncurry f) j bs
diff --git a/visualisers/hmvis/src/hmvis/annotated.cljs b/visualisers/hmvis/src/hmvis/annotated.cljs
index 3ab9d40..ed89741 100644
--- a/visualisers/hmvis/src/hmvis/annotated.cljs
+++ b/visualisers/hmvis/src/hmvis/annotated.cljs
@@ -103,6 +103,15 @@
 (defn LitExpr [_ l]
   [:code (str l)])
 
+(defn Alter [colours a]
+  (pprint a)
+  [:code "<alter>"])
+
+(defn CaseExpr [colours e as]
+  [:<> "case " [Expr colours 0 e] " of { "
+       "<alters>"
+       " }"])
+
 (defn Expr [[c & colours] p {e :e t :type}]
   (match e
     {:InL {:tag "LamF" :contents [bs body & _]}}
@@ -118,6 +127,9 @@
                     [Typed c t [LetExpr colours r bs body]])
     {:InL {:tag "LitF" :contents l}}
       [Typed c t [LitExpr colours l]]
+    {:InR {:tag "CaseEF" :contents [scrut as]}}
+      (maybe-parens (< ppr/app-prec1 p)
+                    [Typed c t [CaseExpr colours scrut as]])
     :else [:code "<expr>"]))
 
 (def rainbow-cycle (cycle ["red"