adt support in type inference

README.org

@@ -101,7 +101,8 @@ For the time being, I just disabled the memoisation. This is very, very bad.
 
 ** TODO [#A] ~case~ inference                                          :feature:
 
-** TODO [#A] ADT support in Rlp/HindleyMilner.hs                       :feature:
+** DONE [#A] ADT support in Rlp/HindleyMilner.hs                       :feature:
+   CLOSED: [2024-04-05 Fri 12:28]
 
 ** DONE whole-program inference (wrap top-level in a ~letrec~)         :feature:
    CLOSED: [2024-04-04 Thu 12:42]

src/Rlp/AltSyntax.hs

@@ -2,7 +2,7 @@
 module Rlp.AltSyntax
     (
     -- * AST
-      Program(..), Decl(..), ExprF(..), Pat(..)
+      Program(..), Decl(..), ExprF(..), Pat(..), pattern ConP'
     , RlpExprF, RlpExpr, Binding(..), Alter(..)
     , RlpExpr', RlpExprF', AnnotatedRlpExpr', Type'
     , DataCon(..), Type(..), Kind
@@ -18,7 +18,7 @@ module Rlp.AltSyntax
     -- * Optics
     , programDecls
     , _VarP, _FunB, _VarB
-    , _TySigD, _FunD
+    , _TySigD, _FunD, _DataD
     , _LetEF
     , Core.applicants1, Core.arrowStops
 
@@ -141,6 +141,20 @@ data Pat b = VarP b
            | AppP (Pat b) (Pat b)
            deriving (Eq, Show, Generic, Generic1)
 
+conList :: Prism' (Pat b) (b, [Pat b])
+conList = prism' up down where
+    up (b,as)               = foldl AppP (ConP b) as
+    down (ConP b)           = Just (b, [])
+    down (AppP (ConP b) as) = Just (b, go as)
+    down _                  = Nothing
+
+    go (AppP f x) = f : go x
+    go p          = [p]
+
+pattern ConP' :: b -> [Pat b] -> Pat b
+pattern ConP' c as <- (preview conList -> Just (c,as))
+    where ConP' c as = review conList (c,as)
+
 deriveShow1 ''Alter
 deriveShow1 ''Binding
 deriveShow1 ''ExprF

src/Rlp/HindleyMilner.hs

@@ -136,6 +136,11 @@ 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
+    pure (ta,j)
+
 forBinds :: (PsName -> Type' -> Judgement -> Judgement)
          -> [(PsName, Type')] -> Judgement -> Judgement
 forBinds f bs j = foldr (uncurry f) j bs
@@ -165,6 +170,9 @@ unify (c:cs) = case c of
     Equality (s :-> t) (s' :-> t')
         -> unify (Equality s s' : Equality t t' : cs)
 
+    Equality (AppT s t) (AppT s' t')
+        -> unify (Equality s s' : Equality t t' : cs)
+
     ImplicitInstance m s t
         | null $ (freeTvs t `S.difference` freeTvs m)
                     `S.intersection` activeTvs cs
@@ -198,11 +206,19 @@ generalise m t = foldr ForallT t as
 occurs :: (HasTypes a) => Name -> a -> Bool
 occurs x t = x `elem` freeTvs t
 
+elimGlobalBinds :: [(Name, Scheme)] -> Cofree RlpExprF' (Type', Judgement)
+                -> Cofree RlpExprF' (Type', Judgement)
+elimGlobalBinds bs = traversed . _2 %~ forBinds f bs where
+    f n t@(ForallT _ _) = elimGenerally n t
+    f n t               = elim n t
+
 --------------------------------------------------------------------------------
 
 annotate :: (Unique :> es)
          => RlpExpr' -> Eff es (Cofree RlpExprF' (Type', Judgement))
-annotate = dendroscribeM (gather . fmap extract)
+annotate = fmap (elimGlobalBinds [ ("Just", ForallT "a" $ VarT "a" :-> ConT "Maybe" `AppT` VarT "a")
+                                 , ("isJust", ForallT "a" $ ConT "Maybe" `AppT` VarT "a" :-> ConT "Bool")])
+         . dendroscribeM (gather . fmap extract)
 
 orderConstraints :: [Constraint] -> [Constraint]
 orderConstraints cs = a <> b
@@ -266,21 +282,32 @@ annotateDefs :: (Unique :> es)
                          (Cofree RlpExprF' (Type', Judgement)))
 annotateDefs = traverseOf (programDefs . _2) annotate
 
+extractDefs :: Program PsName (Cofree RlpExprF' (Type', Judgement))
+            -> [(Name, Type')]
+extractDefs p = p ^.. programDefs & each . _2 %~ fst . extract
+
+extractCons :: Program PsName (Cofree RlpExprF' (Type', Judgement))
+            -> [(Name, Type')]
+extractCons = foldMapOf (programDecls . each . _DataD) \(n,as,cs) ->
+    let root = foldl AppT (ConT n) (VarT <$> as)
+    in cs & fmap \ (DataCon cn cas) -> (cn, foldr (:->) root cas)
+
 annotateProg :: (Unique :> es)
              => Program PsName RlpExpr'
              -> Eff es (Program PsName
                          (Cofree RlpExprF' (Type', Judgement)))
 annotateProg p = do
     p' <- annotateDefs p
-    let bs = p' ^.. programDefs & each . _2 %~ (fst . extract)
+    let bs = extractCons p' ++ extractDefs p'
         p'' = p' & programDefs . _2 . traversed . _2
             %~ forBinds elimGenerally bs
     pure p''
 
 programDefs :: Traversal (Program b a) (Program b a') (b, a) (b, a')
-programDefs k (Program ds) = Program <$> go k ds where
+programDefs k (Program ds) = Program <$> traverse go ds where
-    go k [] = pure []
+    go (FunD n as e) = refun as (k (n,e))
-    go k (FunD n as e : ds) = (:) <$> refun as (k (n,e)) <*> go k ds
+    go (DataD n as cs) = pure $ DataD n as cs
+    go (TySigD n ts) = pure $ TySigD n ts
     refun as kne = uncurry (\a b -> FunD a as b) <$> kne
 
 --------------------------------------------------------------------------------