update: refactor complete

Author: cstml

lambda-buffers-compiler/src/LambdaBuffers/Compiler/KindCheck.hs

@@ -1,8 +1,13 @@
 module LambdaBuffers.Compiler.KindCheck.Derivation (
-  Derivation (Axiom, Abstraction, Application),
+  Derivation (Axiom, Abstraction, Application, Implication),
+  dType,
+  dTopKind,
 ) where
 
-import LambdaBuffers.Compiler.KindCheck.Judgement (Judgement)
+import Control.Lens (Lens', lens, (&), (.~), (^.))
+import LambdaBuffers.Compiler.KindCheck.Judgement (Judgement, jKind, jType)
+import LambdaBuffers.Compiler.KindCheck.Kind (Kind)
+import LambdaBuffers.Compiler.KindCheck.Type (Type)
 import Prettyprinter (
   Doc,
   Pretty (pretty),
@@ -18,13 +23,45 @@ data Derivation
   = Axiom Judgement
   | Abstraction Judgement Derivation
   | Application Judgement Derivation Derivation
+  | Implication Judgement Derivation
   deriving stock (Show, Eq)
 
 instance Pretty Derivation where
   pretty x = case x of
     Axiom j -> hang 2 $ pretty j
     Abstraction j d -> dNest j [d]
     Application j d1 d2 -> dNest j [d1, d2]
+    Implication j d -> dNest j [d]
     where
       dNest :: forall a b c. (Pretty a, Pretty b) => a -> [b] -> Doc c
       dNest j ds = pretty j <> line <> hang 2 (encloseSep (lbracket <> space) rbracket (space <> "∧" <> space) (pretty <$> ds))
+
+dType :: Lens' Derivation Type
+dType = lens from to
+  where
+    from = \case
+      Axiom j -> j ^. jType
+      Abstraction j _ -> j ^. jType
+      Application j _ _ -> j ^. jType
+      Implication j _ -> j ^. jType
+
+    to drv t = case drv of
+      Axiom j -> Axiom $ j & jType .~ t
+      Abstraction j d -> Abstraction (j & jType .~ t) d
+      Application j d1 d2 -> Application (j & jType .~ t) d1 d2
+      Implication j d -> Implication (j & jType .~ t) d
+
+dTopKind :: Lens' Derivation Kind
+dTopKind = lens from to
+  where
+    from = \case
+      Axiom j -> j ^. jKind
+      Abstraction j _ -> j ^. jKind
+      Application j _ _ -> j ^. jKind
+      Implication j _ -> j ^. jKind
+
+    to der t = case der of
+      Axiom j -> Axiom $ j & jKind .~ t
+      Abstraction j d -> Abstraction (j & jKind .~ t) d
+      Application j d1 d2 -> Application (j & jKind .~ t) d1 d2
+      Implication j d -> Abstraction (j & jKind .~ t) d

lambda-buffers-compiler/src/LambdaBuffers/Compiler/KindCheck/Derivation.hs

@@ -1,3 +1,4 @@
+{-# OPTIONS_GHC -Wno-missing-local-signatures #-}
 {-# OPTIONS_GHC -Wno-redundant-constraints #-}
 
 -- This pragma^ is needed due to redundant constraint in Getter.
@@ -16,24 +17,27 @@ module LambdaBuffers.Compiler.KindCheck.Inference (
 ) where
 
 import Data.Bifunctor (Bifunctor (second))
+import Data.Foldable (foldrM)
 
 import LambdaBuffers.Compiler.KindCheck.Context (Context (Context), addContext, context, getAllContext)
-import LambdaBuffers.Compiler.KindCheck.Derivation (Derivation (Abstraction, Application, Axiom))
+import LambdaBuffers.Compiler.KindCheck.Derivation (Derivation (Abstraction, Application, Axiom, Implication), dTopKind, dType)
 import LambdaBuffers.Compiler.KindCheck.Judgement (Judgement (Judgement))
-import LambdaBuffers.Compiler.KindCheck.Kind (Kind (KVar, Type, (:->:)))
-import LambdaBuffers.Compiler.KindCheck.Type (Type (Abs, App, Var), tyOpaque, tyProd, tySum, tyUnit, tyVoid)
-import LambdaBuffers.Compiler.KindCheck.Variable (Atom, Variable)
+import LambdaBuffers.Compiler.KindCheck.Kind (Kind (KType, KVar, (:->:)))
+import LambdaBuffers.Compiler.KindCheck.Type (Type (Abs, App, Constructor, Opaque, Product, Sum, Var, VoidT))
+import LambdaBuffers.Compiler.KindCheck.Variable (Atom, Variable (ForeignRef, LocalRef, TyVar))
 
 import Control.Monad.Freer (Eff, Member, Members, run)
 import Control.Monad.Freer.Error (Error, runError, throwError)
-import Control.Monad.Freer.Reader (Reader, ask, asks, local, runReader)
+import Control.Monad.Freer.Reader (Reader, ask, asks, runReader)
 import Control.Monad.Freer.State (State, evalState, get, put)
 import Control.Monad.Freer.Writer (Writer, runWriter, tell)
 
+import LambdaBuffers.Compiler.ProtoCompat qualified as PC
+
 import Data.String (fromString)
 import Data.Text qualified as T
 
-import Control.Lens (Getter, to, (&), (.~), (^.))
+import Control.Lens ((&), (.~), (^.))
 import Data.Map qualified as M
 
 import Prettyprinter (
@@ -71,6 +75,7 @@ type Derive a =
   forall effs.
   Members
     '[ Reader Context
+     , Reader Kind
      , State DerivationContext
      , Writer [Constraint]
      , Error InferErr
@@ -82,55 +87,26 @@ type Derive a =
 -- Runners
 
 -- | Run derivation builder - not unified yet.
-runDerive :: Context -> Type -> Either InferErr (Derivation, [Constraint])
-runDerive ctx t = run $ runError $ runWriter $ evalState (DC atoms) $ runReader ctx (derive t)
+runDerive :: Context -> PC.TyAbs -> Kind -> Either InferErr (Derivation, [Constraint])
+runDerive ctx t k = run $ runError $ runWriter $ evalState (DC atoms) $ runReader ctx $ runReader k (derive t)
 
-infer :: Context -> Type -> Either InferErr Kind
-infer ctx t = do
-  (d, c) <- runDerive (defContext <> ctx) t
+infer :: Context -> PC.TyDef -> Kind -> Either InferErr Kind
+infer ctx t k = do
+  (d, c) <- runDerive (defContext <> ctx) (t ^. #tyAbs) k
   s <- runUnify' c
   let res = foldl (flip substitute) d s
-  pure $ res ^. topKind
+  pure $ res ^. dTopKind
 
 -- | Default KC Context.
 defContext :: Context
-defContext =
-  mempty
-    & context
-      .~ M.fromList
-        [ (tySum, Type :->: Type :->: Type)
-        , (tyProd, Type :->: Type :->: Type)
-        , (tyUnit, Type)
-        , (tyVoid, Type)
-        , (tyOpaque, Type)
-        ]
+defContext = mempty
 
 --------------------------------------------------------------------------------
 -- Implementation
 
 -- | Creates the derivation
-derive :: Type -> Derive Derivation
-derive x = do
-  c <- ask
-  case x of
-    Var at -> do
-      v <- getBinding at
-      pure $ Axiom $ Judgement (c, x, v)
-    App t1 t2 -> do
-      d1 <- derive t1
-      d2 <- derive t2
-      let ty1 = d1 ^. topKind
-          ty2 = d2 ^. topKind
-      v <- KVar <$> fresh
-      tell [Constraint (ty1, ty2 :->: v)]
-      pure $ Application (Judgement (c, x, v)) d1 d2
-    Abs v t -> do
-      newTy <- getBinding v
-      d <- local (\(Context ctx addC) -> Context ctx $ M.insert v newTy addC) (derive t)
-      let ty = d ^. topKind
-      freshT <- KVar <$> fresh
-      tell [Constraint (freshT, newTy :->: ty)]
-      pure $ Abstraction (Judgement (c, x, freshT)) d
+derive :: PC.TyAbs -> Derive Derivation
+derive x = deriveTyAbs x
   where
     fresh :: Derive Atom
     fresh = do
@@ -139,6 +115,130 @@ derive x = do
         a : as -> put (DC as) >> pure a
         [] -> throwError $ InferImpossibleErr "Reached end of infinite stream."
 
+    deriveTyAbs :: PC.TyAbs -> Derive Derivation
+    deriveTyAbs tyabs =
+      case M.toList (tyabs ^. #tyArgs) of
+        [] -> deriveTyBody (x ^. #tyBody)
+        a@(n, _) : as -> do
+          vK <- getBinding (TyVar n)
+          freshT <- KVar <$> fresh
+          let newAbs = tyabs & #tyArgs .~ uncurry M.singleton a
+          let restAbs = tyabs & #tyArgs .~ M.fromList as
+          restF <- deriveTyAbs restAbs
+          let uK = restF ^. dTopKind
+          tell [Constraint (freshT, uK)]
+          ctx <- ask
+          pure $ Abstraction (Judgement (ctx, Abs newAbs, vK :->: freshT)) restF
+
+    deriveTyBody :: PC.TyBody -> Derive Derivation
+    deriveTyBody = \case
+      PC.OpaqueI si -> do
+        ctx <- ask
+        pure $ Axiom $ Judgement (ctx, Opaque si, KType)
+      PC.SumI s -> deriveSum s
+
+    deriveSum :: PC.Sum -> Derive Derivation
+    deriveSum s = do
+      case M.toList (s ^. #constructors) of
+        [] -> voidDerivation
+        c : cs -> do
+          dc <- deriveConstructor $ snd c
+          restDc <- deriveSum $ s & #constructors .~ M.fromList cs
+          sumDerivation dc restDc
+
+    deriveConstructor :: PC.Constructor -> Derive Derivation
+    deriveConstructor c = do
+      ctx <- ask
+      d <- deriveProduct (c ^. #product)
+      tell $ Constraint <$> [(KType, d ^. dTopKind)]
+      pure $ Implication (Judgement (ctx, Constructor c, d ^. dTopKind)) d
+
+    deriveProduct :: PC.Product -> Derive Derivation
+    deriveProduct = \case
+      PC.RecordI r -> deriveRecord r
+      PC.TupleI t -> deriveTuple t
+
+    deriveRecord r = do
+      case M.toList (r ^. #fields) of
+        [] -> voidDerivation
+        f : fs -> do
+          d1 <- deriveField $ snd f
+          d2 <- deriveRecord $ r & #fields .~ M.fromList fs
+          productDerivation d1 d2
+
+    deriveField :: PC.Field -> Derive Derivation
+    deriveField f = deriveTy $ f ^. #fieldTy
+
+    deriveTy :: PC.Ty -> Derive Derivation
+    deriveTy = \case
+      PC.TyVarI tv -> deriveTyVar tv
+      PC.TyAppI ta -> deriveTyApp ta
+      PC.TyRefI tr -> deriveTyRef tr
+
+    deriveTyRef :: PC.TyRef -> Derive Derivation
+    deriveTyRef = \case
+      PC.LocalI r -> do
+        let ty = LocalRef r
+        v <- getBinding ty
+        c <- ask
+        pure . Axiom . Judgement $ (c, Var ty, v)
+      PC.ForeignI r -> do
+        let ty = ForeignRef r
+        v <- getBinding ty
+        c <- ask
+        pure . Axiom . Judgement $ (c, Var ty, v)
+
+    deriveTyVar :: PC.TyVar -> Derive Derivation
+    deriveTyVar tv = do
+      let varName = tv ^. #varName
+      v <- getBinding $ TyVar varName
+      c <- ask
+      pure . Axiom . Judgement $ (c, Var $ TyVar varName, v)
+
+    deriveTyApp :: PC.TyApp -> Derive Derivation
+    deriveTyApp ap = do
+      f <- deriveTy (ap ^. #tyFunc)
+      args <- deriveTy `traverse` (ap ^. #tyArgs)
+      applyDerivation $ f : args
+
+    deriveTuple :: PC.Tuple -> Derive Derivation
+    deriveTuple t = do
+      voidD <- voidDerivation
+      ds <- deriveTy `traverse` (t ^. #fields)
+      foldrM productDerivation voidD ds
+
+    voidDerivation :: Derive Derivation
+    voidDerivation = do
+      ctx <- ask
+      pure $ Axiom $ Judgement (ctx, VoidT, KType)
+
+    productDerivation :: Derivation -> Derivation -> Derive Derivation
+    productDerivation d1 d2 = do
+      ctx <- ask
+      let t1 = d1 ^. dType
+      let t2 = d2 ^. dType
+      tell $ Constraint <$> [(d1 ^. dTopKind, KType), (d2 ^. dTopKind, KType)]
+      pure $ Application (Judgement (ctx, Product t1 t2, KType)) d1 d2
+
+    sumDerivation :: Derivation -> Derivation -> Derive Derivation
+    sumDerivation d1 d2 = do
+      ctx <- ask
+      let t1 = d1 ^. dType
+      let t2 = d2 ^. dType
+      tell $ Constraint <$> [(d1 ^. dTopKind, KType), (d2 ^. dTopKind, KType)]
+      pure $ Application (Judgement (ctx, Sum t1 t2, KType)) d1 d2
+
+    applyDerivation :: [Derivation] -> Derive Derivation
+    applyDerivation = \case
+      [] -> error "Impossible"
+      [y] -> pure y
+      d1 : ys -> do
+        c <- ask
+        d2 <- applyDerivation ys
+        v <- KVar <$> fresh
+        tell [Constraint ((d2 ^. dTopKind) :->: v, d1 ^. dTopKind)]
+        pure $ Application (Judgement (c, App (d1 ^. dType) (d2 ^. dType), v)) d1 d2
+
 {- | Gets the binding from the context - if the variable is not bound throw an
  error.
 -}
@@ -149,22 +249,13 @@ getBinding t = do
     Just x -> pure x
     Nothing -> throwError $ InferUnboundTermErr t
 
--- | Gets kind from a derivation.
-topKind :: Getter Derivation Kind
-topKind = to f
-  where
-    f = \case
-      Axiom (Judgement (_, _, k)) -> k
-      Abstraction (Judgement (_, _, k)) _ -> k
-      Application (Judgement (_, _, k)) _ _ -> k
-
 -- | Unification monad.
 type Unifier a = forall effs. Member (Error InferErr) effs => Eff effs a
 
 -- | Gets the variables of a type.
 getVariables :: Kind -> [Atom]
 getVariables = \case
-  Type -> mempty
+  KType -> mempty
   x :->: y -> getVariables x <> getVariables y
   KVar x -> [x]
 
@@ -175,14 +266,14 @@ getVariables = \case
 unify :: [Constraint] -> Unifier [Substitution]
 unify [] = pure []
 unify (constraint@(Constraint (l, r)) : xs) = case l of
-  Type -> case r of
-    Type -> unify xs
+  KType -> case r of
+    KType -> unify xs
     (_ :->: _) -> nope constraint
     KVar v ->
-      let sub = Substitution (v, Type)
+      let sub = Substitution (v, KType)
        in (sub :) <$> unify (sub `substituteIn` xs)
   x :->: y -> case r of
-    Type -> nope constraint
+    KType -> nope constraint
     KVar v ->
       if v `appearsIn` l
         then appearsErr v l
@@ -229,7 +320,7 @@ unify (constraint@(Constraint (l, r)) : xs) = case l of
 -- | Applies substitutions to a kind.
 applySubstitution :: Substitution -> Kind -> Kind
 applySubstitution s@(Substitution (a, t)) k = case k of
-  Type -> Type
+  KType -> KType
   l :->: r -> applySubstitution s l :->: applySubstitution s r
   KVar v -> if v == a then t else k
 
@@ -245,6 +336,7 @@ substitute s d = case d of
   Axiom j -> Axiom (applySubsToJudgement s j)
   Abstraction j dc -> Abstraction (applySubsToJudgement s j) (substitute s dc)
   Application j d1 d2 -> Application (applySubsToJudgement s j) (substitute s d1) (substitute s d2)
+  Implication j dc -> Implication (applySubsToJudgement s j) (substitute s dc)
   where
     applySubsToJudgement sub (Judgement (ctx, t, k)) = Judgement (applySubstitutionCtx s ctx, t, applySubstitution sub k)
 

lambda-buffers-compiler/src/LambdaBuffers/Compiler/KindCheck/Inference.hs

@@ -1,5 +1,6 @@
-module LambdaBuffers.Compiler.KindCheck.Judgement (Judgement (Judgement), getJudgement) where
+module LambdaBuffers.Compiler.KindCheck.Judgement (Judgement (Judgement), getJudgement, jType, jKind) where
 
+import Control.Lens (Lens', lens)
 import LambdaBuffers.Compiler.KindCheck.Context (Context)
 import LambdaBuffers.Compiler.KindCheck.Kind (Kind)
 import LambdaBuffers.Compiler.KindCheck.Type (Type)
@@ -8,5 +9,17 @@ import Prettyprinter (Pretty (pretty), (<+>))
 newtype Judgement = Judgement {getJudgement :: (Context, Type, Kind)}
   deriving stock (Show, Eq)
 
+jType :: Lens' Judgement Type
+jType = lens from to
+  where
+    from = (\(_, x, _) -> x) . getJudgement
+    to (Judgement (c, _, k)) t = Judgement (c, t, k)
+
+jKind :: Lens' Judgement Kind
+jKind = lens from to
+  where
+    from = (\(_, _, k) -> k) . getJudgement
+    to (Judgement (c, t, _)) k = Judgement (c, t, k)
+
 instance Pretty Judgement where
   pretty (Judgement (c, t, k)) = pretty c <> " ⊢ " <> pretty t <+> ":" <+> pretty k