[microjuvix] implement basic typechecker

app/Main.hs

@@ -20,6 +20,7 @@ import MiniJuvix.Syntax.Concrete.Scoped.Pretty.Html
 import qualified MiniJuvix.Syntax.Concrete.Scoped.Pretty.Text as T
 import qualified MiniJuvix.Syntax.Concrete.Scoped.Scoper as M
 import qualified MiniJuvix.Syntax.MicroJuvix.Pretty.Ansi as Micro
+import qualified MiniJuvix.Syntax.MicroJuvix.TypeChecker as Micro
 import qualified MiniJuvix.Termination as T
 import qualified MiniJuvix.Termination.CallGraph as A
 import qualified MiniJuvix.Translation.AbstractToMicroJuvix as Micro
@@ -255,8 +256,12 @@ go c = do
       m <- parseModuleIO _mjuvixInputFile
       (_, s) <- fromRightIO' printErrorAnsi $ M.scopeCheck1IO root m
       a <- fromRightIO' putStrLn (return $ A.translateModule s)
-      let mini = Micro.translateModule a
-      Micro.printPrettyCodeDefault mini
+      let micro = Micro.translateModule a
+      Micro.printPrettyCodeDefault micro
+      putStrLn ""
+      case Micro.checkModule micro of
+        Left er -> putStrLn er
+        Right {} -> putStrLn "Well done! It type checks"
     MiniHaskell MiniHaskellOptions {..} -> do
       m <- parseModuleIO _mhaskellInputFile
       (_ , s) <- fromRightIO' printErrorAnsi $ M.scopeCheck1IO root m

src/MiniJuvix/Syntax/MicroJuvix/Language.hs

@@ -45,11 +45,14 @@ data Module = Module
   }
 
 data ModuleBody = ModuleBody
-  { _moduleInductives :: HashMap InductiveName (Indexed InductiveDef),
-    _moduleFunctions :: HashMap FunctionName (Indexed FunctionDef),
-    _moduleForeigns :: [Indexed ForeignBlock]
+  { _moduleStatements :: [Statement]
   }
 
+data Statement =
+  StatementInductive InductiveDef
+  | StatementFunction FunctionDef
+  | StatementForeign ForeignBlock
+
 data FunctionDef = FunctionDef
   { _funDefName :: FunctionName,
     _funDefTypeSig :: Type,
@@ -66,9 +69,15 @@ data Iden
   | IdenConstructor Name
   | IdenVar VarName
 
+data TypedExpression = TypedExpression {
+  _typedType :: Type,
+  _typedExpression :: Expression
+  }
+
 data Expression
   = ExpressionIden Iden
   | ExpressionApplication Application
+  | ExpressionTyped TypedExpression
 
 data Application = Application
   { _appLeft :: Expression,
@@ -79,6 +88,7 @@ data Function = Function
   { _funLeft :: Type,
     _funRight :: Type
   }
+  deriving stock (Eq)
 
 -- | Fully applied constructor in a pattern.
 data ConstructorApp = ConstructorApp
@@ -103,37 +113,40 @@ data InductiveConstructorDef = InductiveConstructorDef
 
 newtype TypeIden
   = TypeIdenInductive InductiveName
+  deriving stock (Eq)
 
 data Type
   = TypeIden TypeIden
   | TypeFunction Function
+  deriving stock (Eq)
+
+data ConstructorInfo = ConstructorInfo {
+  _constructorInfoArgs :: [Type],
+  _constructorInfoInductive :: InductiveName
+  }
+
+data FunctionInfo = FunctionInfo {
+  _functionInfoType :: Type
+  }
+
+data InfoTable = InfoTable {
+  _infoConstructors :: HashMap Name ConstructorInfo,
+  _infoFunctions :: HashMap Name FunctionInfo
+  }
 
 makeLenses ''Module
 makeLenses ''Function
 makeLenses ''FunctionDef
+makeLenses ''FunctionInfo
+makeLenses ''ConstructorInfo
 makeLenses ''FunctionClause
 makeLenses ''InductiveDef
 makeLenses ''ModuleBody
 makeLenses ''Application
+makeLenses ''TypedExpression
 makeLenses ''InductiveConstructorDef
 makeLenses ''ConstructorApp
 
-instance Semigroup ModuleBody where
-  a <> b =
-    ModuleBody
-      { _moduleInductives = a ^. moduleInductives <> b ^. moduleInductives,
-        _moduleFunctions = a ^. moduleFunctions <> b ^. moduleFunctions,
-        _moduleForeigns = a ^. moduleForeigns <> b ^. moduleForeigns
-      }
-
-instance Monoid ModuleBody where
-  mempty =
-    ModuleBody
-      { _moduleInductives = mempty,
-        _moduleForeigns = mempty,
-        _moduleFunctions = mempty
-      }
-
 instance HasAtomicity Application where
   atomicity = const (Aggregate appFixity)
 
@@ -141,6 +154,7 @@ instance HasAtomicity Expression where
   atomicity e = case e of
     ExpressionIden {} -> Atom
     ExpressionApplication a -> atomicity a
+    ExpressionTyped t -> atomicity (t ^. typedExpression)
 
 instance HasAtomicity Function where
   atomicity = const (Aggregate funFixity)

src/MiniJuvix/Syntax/MicroJuvix/Pretty/Base.hs

@@ -41,10 +41,14 @@ instance PrettyCode Application where
     r' <- ppRightExpression appFixity (a ^. appRight)
     return $ l' <+> r'
 
+instance PrettyCode TypedExpression where
+  ppCode e = ppCode (e ^. typedExpression)
+
 instance PrettyCode Expression where
   ppCode e = case e of
     ExpressionIden i -> ppCode i
     ExpressionApplication a -> ppCode a
+    ExpressionTyped a -> ppCode a
 
 keyword :: Text -> Doc Ann
 keyword = annotate AnnKeyword . pretty
@@ -152,13 +156,15 @@ instance PrettyCode ForeignBlock where
         <> line
         <> rbrace
 
--- TODO Jonathan review
+instance PrettyCode Statement where
+  ppCode = \case
+    StatementForeign f -> ppCode f
+    StatementFunction f -> ppCode f
+    StatementInductive f -> ppCode f
+
 instance PrettyCode ModuleBody where
   ppCode m = do
-    types' <- mapM (mapM ppCode) (toList (m ^. moduleInductives))
-    funs' <- mapM (mapM ppCode) (toList (m ^. moduleFunctions))
-    foreigns' <- mapM (mapM ppCode) (toList (m ^. moduleForeigns))
-    let everything = map (^. indexedThing) (sortOn (^. indexedIx) (types' ++ funs' ++ foreigns'))
+    everything <- mapM ppCode (m ^. moduleStatements)
     return $ vsep2 everything
     where
       vsep2 = concatWith (\a b -> a <> line <> line <> b)

src/MiniJuvix/Syntax/MicroJuvix/TypeChecker.hs

@@ -0,0 +1,172 @@
+module MiniJuvix.Syntax.MicroJuvix.TypeChecker where
+import MiniJuvix.Prelude
+import MiniJuvix.Syntax.MicroJuvix.Language
+import qualified Data.HashMap.Strict as HashMap
+
+type Err = Text
+
+newtype LocalVars = LocalVars {
+  _localTypes :: HashMap VarName Type
+  }
+  deriving newtype (Semigroup, Monoid)
+makeLenses ''LocalVars
+
+checkModule :: Module -> Either Err Module
+checkModule m = run $ runError $ runReader (buildTable m) (checkModule' m)
+
+buildTable :: Module -> InfoTable
+buildTable m = InfoTable {..}
+  where
+   _infoConstructors :: HashMap Name ConstructorInfo
+   _infoConstructors = HashMap.fromList
+     [ (c ^. constructorName, ConstructorInfo args ind) |
+       StatementInductive d <- ss,
+       let ind = d ^. inductiveName,
+       c <- d ^. inductiveConstructors,
+       let args = c ^. constructorParameters
+     ]
+   _infoFunctions :: HashMap Name FunctionInfo
+   _infoFunctions = HashMap.fromList
+     [ (f ^. funDefName, FunctionInfo (f ^. funDefTypeSig)) |
+       StatementFunction f <- ss]
+   ss = m ^. moduleBody ^. moduleStatements
+
+checkModule' :: Members '[Reader InfoTable, Error Err] r =>
+  Module -> Sem r Module
+checkModule' Module {..} = do
+  _moduleBody' <- checkModuleBody _moduleBody
+  return Module {
+    _moduleBody = _moduleBody',
+    ..
+    }
+
+checkModuleBody :: Members '[Reader InfoTable, Error Err] r =>
+  ModuleBody -> Sem r ModuleBody
+checkModuleBody ModuleBody {..} = do
+  _moduleStatements' <- mapM checkStatement _moduleStatements
+  return ModuleBody {
+    _moduleStatements = _moduleStatements'
+    }
+
+checkStatement :: Members '[Reader InfoTable, Error Err] r =>
+  Statement -> Sem r Statement
+checkStatement s = case s of
+  StatementFunction fun -> StatementFunction <$> checkFunctionDef fun
+  StatementForeign {} -> return s
+  StatementInductive {} -> return s -- TODO is checking inductives needed?
+
+checkFunctionDef :: Members '[Reader InfoTable, Error Err] r =>
+  FunctionDef -> Sem r FunctionDef
+checkFunctionDef FunctionDef {..} = do
+  info <- lookupFunction _funDefName
+  _funDefClauses' <- mapM (checkFunctionClause info) _funDefClauses
+  return FunctionDef {
+    _funDefClauses = _funDefClauses',
+    ..
+    }
+
+checkExpression :: Members '[Reader InfoTable, Error Err, Reader LocalVars] r =>
+  Type -> Expression -> Sem r Expression
+checkExpression t e = do
+  t' <- inferExpression' e
+  when (t /= t' ^. typedType) (throwErr "wrong type")
+  return (ExpressionTyped t')
+
+inferExpression :: Members '[Reader InfoTable, Error Err, Reader LocalVars] r =>
+   Expression -> Sem r Expression
+inferExpression = fmap ExpressionTyped . inferExpression'
+
+lookupConstructor :: Member (Reader InfoTable) r => Name -> Sem r ConstructorInfo
+lookupConstructor f = HashMap.lookupDefault impossible f <$> asks _infoConstructors
+
+lookupFunction :: Member (Reader InfoTable) r => Name -> Sem r FunctionInfo
+lookupFunction f = HashMap.lookupDefault impossible f <$> asks _infoFunctions
+
+lookupVar :: Member (Reader LocalVars) r => Name -> Sem r Type
+lookupVar v = HashMap.lookupDefault impossible v <$> asks _localTypes
+
+constructorType :: Member (Reader InfoTable) r => Name -> Sem r Type
+constructorType c = do
+  info <- lookupConstructor c
+  let r = TypeIden (TypeIdenInductive (info ^. constructorInfoInductive))
+  return (foldFunType (info ^. constructorInfoArgs) r)
+
+-- | [a, b] c ==> a -> (b -> c)
+foldFunType :: [Type] -> Type -> Type
+foldFunType l r = case l of
+  [] -> r
+  (a : as) -> TypeFunction (Function a (foldFunType as r))
+
+-- |  a -> (b -> c)  ==> ([a, b], c)
+unfoldFunType :: Type -> ([Type], Type)
+unfoldFunType t = case t of
+  TypeIden {} -> ([], t)
+  TypeFunction (Function l r) -> first (l:) (unfoldFunType r)
+
+throwErr :: Members '[Error Err] r => Err -> Sem r a
+throwErr = throw
+
+inferExpression' :: forall r. Members '[Reader InfoTable, Error Err, Reader LocalVars] r =>
+   Expression -> Sem r TypedExpression
+inferExpression' e = case e of
+  ExpressionIden i -> checkIden i
+  ExpressionApplication a -> checkApplication a
+  ExpressionTyped {} -> impossible
+  where
+  checkIden :: Iden -> Sem r TypedExpression
+  checkIden i = case i of
+    IdenFunction fun -> do
+      info <- lookupFunction fun
+      return (TypedExpression (info ^. functionInfoType) (ExpressionIden i))
+    IdenConstructor c -> do
+      ty <- constructorType c
+      return (TypedExpression ty (ExpressionIden i))
+    IdenVar v -> do
+      ty <- lookupVar v
+      return (TypedExpression ty (ExpressionIden i))
+  checkApplication :: Application -> Sem r TypedExpression
+  checkApplication a = do
+    l <- inferExpression' (a ^. appLeft)
+    fun <- getFunctionType (l ^. typedType)
+    r <- checkExpression (fun ^. funLeft) (a ^. appRight)
+    return TypedExpression {
+      _typedExpression = ExpressionApplication Application {
+          _appLeft = ExpressionTyped l,
+          _appRight = r
+          },
+      _typedType = fun ^. funRight
+      }
+  getFunctionType :: Type -> Sem r Function
+  getFunctionType t = case t of
+    TypeFunction f -> return f
+    _ -> throwErr "expected function type"
+
+checkFunctionClause :: forall r. Members '[Reader InfoTable, Error Err] r =>
+  FunctionInfo -> FunctionClause -> Sem r FunctionClause
+checkFunctionClause info FunctionClause{..} = do
+  let (argTys, rty) = unfoldFunType (info ^. functionInfoType)
+      (patTys, restTys) = splitAt (length _clausePatterns) argTys
+      bodyTy = foldFunType restTys rty
+  when (length patTys /= length _clausePatterns) (throwErr "wrong number of patterns")
+  locals <- mconcat <$> zipWithM checkPattern patTys _clausePatterns
+  clauseBody' <- runReader locals (checkExpression bodyTy _clauseBody)
+  return FunctionClause {
+    _clauseBody = clauseBody',
+    ..
+    }
+
+checkPattern :: forall r. Members '[Reader InfoTable, Error Err] r =>
+  Type -> Pattern -> Sem r LocalVars
+checkPattern type_ pat = LocalVars . HashMap.fromList <$> go type_ pat
+  where
+  go :: Type -> Pattern -> Sem r [(VarName, Type)]
+  go ty p = case p of
+    PatternWildcard -> return []
+    PatternVariable v -> return [(v, ty)]
+    PatternConstructorApp a -> goConstr a
+    where
+    goConstr :: ConstructorApp -> Sem r [(VarName, Type)]
+    goConstr (ConstructorApp c ps) = do
+      tys <- (^. constructorInfoArgs) <$> lookupConstructor c
+      when (length tys /= length ps) (throwErr "wrong number of arguments in constructor app")
+      concat <$> zipWithM go tys ps

src/MiniJuvix/Translation/AbstractToMicroJuvix.hs

@@ -2,6 +2,7 @@ module MiniJuvix.Translation.AbstractToMicroJuvix where
 
 import qualified Data.HashMap.Strict as HashMap
 import MiniJuvix.Prelude
+import MiniJuvix.Syntax.Concrete.Language (ForeignBlock)
 import qualified MiniJuvix.Syntax.Abstract.Language.Extra as A
 import qualified MiniJuvix.Syntax.Concrete.Scoped.Name as S
 import MiniJuvix.Syntax.MicroJuvix.Language
@@ -38,19 +39,23 @@ goModuleBody :: A.ModuleBody -> ModuleBody
 goModuleBody b
   | not (HashMap.null (b ^. A.moduleLocalModules)) = unsupported "local modules"
   | otherwise =
-    ModuleBody
-      { _moduleInductives =
-          HashMap.fromList
-            [ (d ^. indexedThing . inductiveName, d)
-              | d <- map (fmap goInductiveDef) (toList (b ^. A.moduleInductives))
-            ],
-        _moduleFunctions =
-          HashMap.fromList
-            [ (f ^. indexedThing . funDefName, f)
-              | f <- map (fmap goFunctionDef) (toList (b ^. A.moduleFunctions))
-            ],
-        _moduleForeigns = b ^. A.moduleForeigns
-      }
+    ModuleBody sortedStatements
+  where
+   sortedStatements :: [Statement]
+   sortedStatements = map _indexedThing (sortOn _indexedIx statements)
+   statements :: [Indexed Statement]
+   statements = map (fmap StatementForeign) foreigns
+     <> map (fmap StatementFunction) functions
+     <> map (fmap StatementInductive) inductives
+   inductives :: [Indexed InductiveDef]
+   inductives =
+       [ d | d <- map (fmap goInductiveDef) (toList (b ^. A.moduleInductives))]
+   functions :: [Indexed FunctionDef]
+   functions  =
+      [ f | f <- map (fmap goFunctionDef) (toList (b ^. A.moduleFunctions))]
+   foreigns :: [Indexed ForeignBlock]
+   foreigns = b ^. A.moduleForeigns
+
 
 -- <> mconcatMap goImport (b ^. A.moduleImports)