mono fix

src/Monomorphizer/Monomorphizer.hs

@@ -1,71 +1,83 @@
--- | For now, converts polymorphic functions to concrete ones based on usage.
--- Assumes lambdas are lifted.
---
--- This step of compilation is as follows:
---
--- Split all function bindings into monomorphic and polymorphic binds. The
--- monomorphic bindings will be part of this compilation step.
--- Apply the following monomorphization function on all monomorphic binds, with
--- their type as an additional argument.
---
--- The function that transforms Binds operates on both monomorphic and
--- polymorphic functions, creates a context in which all possible polymorphic types
--- are mapped to concrete types, created using the additional argument.
--- Expressions are then recursively processed. The type of these expressions
--- are changed to using the mapped generic types. The expected type provided
--- in the recursion is changed depending on the different nodes.
---
--- When an external bind is encountered (with EId), it is checked whether it
--- exists in outputed binds or not. If it does, nothing further is evaluated.
--- If not, the bind transformer function is called on it with the
--- expected type in this context. The result of this computation (a monomorphic
--- bind) is added to the resulting set of binds.
-
 {-# LANGUAGE LambdaCase #-}
+
+{- | For now, converts polymorphic functions to concrete ones based on usage.
+Assumes lambdas are lifted.
+
+This step of compilation is as follows:
+
+Split all function bindings into monomorphic and polymorphic binds. The
+monomorphic bindings will be part of this compilation step.
+Apply the following monomorphization function on all monomorphic binds, with
+their type as an additional argument.
+
+The function that transforms Binds operates on both monomorphic and
+polymorphic functions, creates a context in which all possible polymorphic types
+are mapped to concrete types, created using the additional argument.
+Expressions are then recursively processed. The type of these expressions
+are changed to using the mapped generic types. The expected type provided
+in the recursion is changed depending on the different nodes.
+
+When an external bind is encountered (with EId), it is checked whether it
+exists in outputed binds or not. If it does, nothing further is evaluated.
+If not, the bind transformer function is called on it with the
+expected type in this context. The result of this computation (a monomorphic
+bind) is added to the resulting set of binds.
+-}
 module Monomorphizer.Monomorphizer (monomorphize, morphExp, morphBind) where
 
 import Monomorphizer.DataTypeRemover (removeDataTypes)
-import qualified Monomorphizer.MonomorphizerIr as O
+import Monomorphizer.MonomorphizerIr qualified as O
-import qualified Monomorphizer.MorbIr          as M
+import Monomorphizer.MorbIr qualified as M
-import qualified TypeChecker.TypeCheckerIr     as T
 import TypeChecker.TypeCheckerIr (Ident (Ident))
+import TypeChecker.TypeCheckerIr qualified as T
 
-import           Control.Monad.Reader          (MonadReader (ask, local),
+import Control.Monad.Reader (
-                                                Reader, asks, runReader, when)
+    MonadReader (ask, local),
-import           Control.Monad.State           (MonadState, StateT (runStateT),
+    Reader,
-                                                gets, modify)
+    asks,
+    runReader,
+    when,
+ )
+import Control.Monad.State (
+    MonadState,
+    StateT (runStateT),
+    gets,
+    modify,
+ )
 import Data.Coerce (coerce)
-import qualified Data.Map                      as Map
+import Data.Map qualified as Map
 import Data.Maybe (fromJust)
-import qualified Data.Set                      as Set
+import Data.Set qualified as Set
 import Debug.Trace
 import Grammar.Print (printTree)
 
--- | EnvM is the monad containing the read-only state as well as the
+{- | EnvM is the monad containing the read-only state as well as the
--- output state containing monomorphized functions and to-be monomorphized
+output state containing monomorphized functions and to-be monomorphized
--- data type declarations.
+data type declarations.
+-}
 newtype EnvM a = EnvM (StateT Output (Reader Env) a)
     deriving (Functor, Applicative, Monad, MonadState Output, MonadReader Env)
 
 type Output = Map.Map Ident Outputted
 
--- | Data structure describing outputted top-level information, that is
+{- | Data structure describing outputted top-level information, that is
--- Binds, Polymorphic Data types (monomorphized in a later step) and
+Binds, Polymorphic Data types (monomorphized in a later step) and
--- Marked bind, which means that it is in the process of monomorphization
+Marked bind, which means that it is in the process of monomorphization
--- and should not be monomorphized again.
+and should not be monomorphized again.
+-}
 data Outputted = Marked | Complete M.Bind | Data M.Type T.Data
 
 -- | Static environment.
-data Env = Env {
+data Env = Env
-  -- | All binds in the program.
+    { input :: Map.Map Ident T.Bind
-  input    :: Map.Map Ident T.Bind,
+    -- ^ All binds in the program.
-  -- | All constructors mapped to their respective polymorphic data def
+    , dataDefs :: Map.Map Ident T.Data
+    -- ^ All constructors mapped to their respective polymorphic data def
     -- which includes all other constructors.
-  dataDefs ::  Map.Map Ident T.Data,
+    , polys :: Map.Map Ident M.Type
-  -- | Maps polymorphic identifiers with concrete types.
+    -- ^ Maps polymorphic identifiers with concrete types.
-  polys    :: Map.Map Ident M.Type,
+    , locals :: Set.Set Ident
-  -- | Local variables.
+    -- ^ Local variables.
-  locals   :: Set.Set Ident
     }
 
 -- | Determines if the identifier describes a local variable in the given context.
@@ -80,8 +92,9 @@ getInputBind ident = asks (Map.lookup ident . input)
 addOutputBind :: M.Bind -> EnvM ()
 addOutputBind b@(M.Bind (ident, _) _ _) = modify (Map.insert ident (Complete b))
 
--- | Marks a global bind as being processed, meaning that when encountered again,
+{- | Marks a global bind as being processed, meaning that when encountered again,
--- it should not be recursively processed.
+it should not be recursively processed.
+-}
 markBind :: Ident -> EnvM ()
 markBind ident = modify (Map.insert ident Marked)
 
@@ -93,21 +106,25 @@ isBindMarked ident = gets (Map.member ident)
 getMain :: EnvM T.Bind
 getMain = asks (\env -> fromJust $ Map.lookup (T.Ident "main") (input env))
 
--- | Makes a kv pair list of polymorphic to monomorphic mappings, throws runtime
+{- | Makes a kv pair list of polymorphic to monomorphic mappings, throws runtime
--- error when encountering different structures between the two arguments.
+error when encountering different structures between the two arguments.
+-}
 mapTypes :: T.Type -> M.Type -> [(Ident, M.Type)]
 mapTypes (T.TLit _) (M.TLit _) = []
 mapTypes (T.TVar (T.MkTVar i1)) tm = [(i1, tm)]
-mapTypes (T.TFun pt1 pt2)        (M.TFun mt1 mt2)  = mapTypes pt1 mt1 ++
+mapTypes (T.TFun pt1 pt2) (M.TFun mt1 mt2) =
-                                                     mapTypes pt2 mt2
+    mapTypes pt1 mt1
-mapTypes (T.TData tIdent pTs)    (M.TData mIdent mTs) = if tIdent /= mIdent
+        ++ mapTypes pt2 mt2
+mapTypes (T.TData tIdent pTs) (M.TData mIdent mTs) =
+    if tIdent /= mIdent
         then error "the data type names of monomorphic and polymorphic data types does not match"
         else foldl (\xs (p, m) -> mapTypes p m ++ xs) [] (zip pTs mTs)
 mapTypes t1 t2 = error $ "structure of types not the same: '" ++ printTree t1 ++ "', '" ++ printTree t2 ++ "'"
 
 -- | Gets the mapped monomorphic type of a polymorphic type in the current context.
 getMonoFromPoly :: T.Type -> EnvM M.Type
-getMonoFromPoly t = do env <- ask
+getMonoFromPoly t = do
+    env <- ask
     return $ getMono (polys env) t
   where
     getMono :: Map.Map Ident M.Type -> T.Type -> M.Type
@@ -120,19 +137,27 @@ getMonoFromPoly t = do env <- ask
         -- error $ "type not found! type: " ++ show ident ++ ", error in previous compilation steps"
         (T.TData ident args) -> M.TData ident (map (getMono polys) args)
 
--- | If ident not already in env's output, morphed bind to output
+{- | If ident not already in env's output, morphed bind to output
--- (and all referenced binds within this bind).
+(and all referenced binds within this bind).
--- Returns the annotated bind name.
+Returns the annotated bind name.
+-}
 morphBind :: M.Type -> T.Bind -> EnvM Ident
 morphBind expectedType b@(T.Bind (Ident str, btype) args (exp, expt)) =
-    local (\env -> env { locals = Set.fromList (map fst args),
+    local
-                         polys  = Map.fromList (mapTypes btype expectedType)
+        ( \env ->
-                       }) $ do
+            env
+                { locals = Set.fromList (map fst args)
+                , polys = Map.fromList (mapTypes btype expectedType)
+                }
+        )
+        $ do
             -- The "new name" is used to find out if it is already marked or not.
             let name' = newFuncName expectedType b
             bindMarked <- isBindMarked (coerce name')
             -- Return with right name if already marked
-      if bindMarked then return name' else do
+            if bindMarked
+                then return name'
+                else do
                     -- Mark so that this bind will not be processed in recursive or cyclic
                     -- function calls
                     markBind (coerce name')
@@ -140,22 +165,28 @@ morphBind expectedType b@(T.Bind (Ident str, btype) args (exp, expt)) =
                     exp' <- morphExp expt' exp
                     -- Get monomorphic type sof args
                     args' <- mapM morphArg args
-        addOutputBind $ M.Bind (coerce name', expectedType)
+                    addOutputBind $
-            args' (exp', expt')
+                        M.Bind
+                            (coerce name', expectedType)
+                            args'
+                            (exp', expt')
                     return name'
 
 -- | Monomorphizes arguments of a bind.
 morphArg :: (Ident, T.Type) -> EnvM (Ident, M.Type)
-morphArg (ident, t) = do t' <- getMonoFromPoly t
+morphArg (ident, t) = do
+    t' <- getMonoFromPoly t
     return (ident, t')
 
 -- | Gets the data bind from the name of a constructor.
 getInputData :: Ident -> EnvM (Maybe T.Data)
-getInputData ident = do env <- ask
+getInputData ident = do
+    env <- ask
     return $ Map.lookup ident (dataDefs env)
 
--- | Monomorphize a constructor using it's global name. Constructors may
+{- | Monomorphize a constructor using it's global name. Constructors may
--- appear as expressions in the tree, or as patterns in case-expressions.
+appear as expressions in the tree, or as patterns in case-expressions.
+-}
 morphCons :: M.Type -> Ident -> EnvM ()
 morphCons expectedType ident = do
     maybeD <- getInputData ident
@@ -197,7 +228,8 @@ morphExp expectedType exp = case exp of
         return $ M.ECase (exp', t') bs'
     T.EVar ident -> do
         isLocal <- localExists ident
-    if isLocal then do
+        if isLocal
+            then do
                 return $ M.EVar (coerce ident)
             else do
                 bind <- getInputBind ident
@@ -210,7 +242,6 @@ morphExp expectedType exp = case exp of
                         -- New bind to process
                         newBindName <- morphBind expectedType bind'
                         return $ M.EVar (coerce newBindName)
-
     T.ELet (T.Bind{}) _ -> error "lets not possible yet"
 
 -- | Monomorphizes case-of branches.
@@ -227,20 +258,24 @@ morphBranch (T.Branch (p, pt) (e, et)) = do
 -- | Morphs pattern (pattern => expression), gives the newly bound local variables.
 morphPattern :: M.Type -> Set.Set Ident -> (T.Pattern, T.Type) -> EnvM (M.Pattern, Set.Set Ident)
 morphPattern expectedType ls (p, t) = case p of
-  T.PVar ident -> do t' <- getMonoFromPoly t
+    T.PVar ident -> do
+        t' <- getMonoFromPoly t
         return (M.PVar (ident, t'), Set.insert ident ls)
-  T.PLit lit ->   do t' <- getMonoFromPoly t
+    T.PLit lit -> do
+        t' <- getMonoFromPoly t
         return (M.PLit (convertLit lit, t'), ls)
     T.PCatch -> return (M.PCatch, ls)
     -- Constructor ident
-  T.PEnum ident   -> do morphCons expectedType ident
+    T.PEnum ident -> do
+        morphCons expectedType ident
         return (M.PEnum ident, ls)
-  T.PInj ident ps -> do morphCons expectedType ident
+    T.PInj ident ps -> do
+        morphCons expectedType ident
         let (M.TData tIdent ts) = expectedType
         -- TODO: this is wrong!
         pairs <- mapM (\(pat, patT) -> morphPattern patT ls pat) (zip ps ts)
-                        if length ts == length ps then
+        if length ts == length ps
-                          return (M.PCatch, Set.singleton $ Ident "$1y")
+            then return (M.PCatch, Set.singleton $ Ident "$1y")
             else return (M.PInj ident (map fst pairs), Set.unions (map snd pairs))
 
 -- | Creates a new identifier for a function with an assigned type.
@@ -260,8 +295,12 @@ newName t (Ident str) = Ident $ str ++ "$" ++ newName' t
 
 -- | Monomorphization step.
 monomorphize :: T.Program -> O.Program
-monomorphize (T.Program defs) = removeDataTypes $ M.Program (getDefsFromOutput
+monomorphize (T.Program defs) =
-    (runEnvM Map.empty (createEnv defs) monomorphize'))
+    removeDataTypes $
+        M.Program
+            ( getDefsFromOutput
+                (runEnvM Map.empty (createEnv defs) monomorphize')
+            )
   where
     monomorphize' :: EnvM ()
     monomorphize' = do
@@ -275,10 +314,13 @@ runEnvM o env (EnvM stateM) = snd $ runReader (runStateT stateM o) env
 
 -- | Creates the environment based on the input binds.
 createEnv :: [T.Def] -> Env
-createEnv defs = Env { input    = Map.fromList bindPairs,
+createEnv defs =
-                       dataDefs = Map.fromList dataPairs,
+    Env
-                       polys    = Map.empty,
+        { input = Map.fromList bindPairs
-                       locals   = Set.empty }
+        , dataDefs = Map.fromList dataPairs
+        , polys = Map.empty
+        , locals = Set.empty
+        }
   where
     bindPairs = (map (\b -> (getBindName b, b)) . getBindsFromDefs) defs
     dataPairs :: [(Ident, T.Data)]
@@ -291,32 +333,42 @@ getBindName (T.Bind (ident, _) _ _) = ident
 -- Helper functions
 -- Gets custom data declarations form defs.
 getDataFromDefs :: [T.Def] -> [T.Data]
-getDataFromDefs = foldl (\bs -> \case
+getDataFromDefs =
+    foldl
+        ( \bs -> \case
             T.DBind _ -> bs
-                                   T.DData d -> d:bs) []
+            T.DData d -> d : bs
+        )
+        []
 
 getConsName :: T.Inj -> Ident
 getConsName (T.Inj ident _) = ident
 
 getBindsFromDefs :: [T.Def] -> [T.Bind]
-getBindsFromDefs = foldl (\bs -> \case
+getBindsFromDefs =
+    foldl
+        ( \bs -> \case
             T.DBind b -> b : bs
-                                   T.DData _ -> bs) []
+            T.DData _ -> bs
+        )
+        []
 
 getDefsFromOutput :: Output -> [M.Def]
 getDefsFromOutput o =
-  map M.DBind binds ++
+    map M.DBind binds
-  (map (M.DData . snd) . Map.toList) (createNewData dataInput Map.empty)
+        ++ (map (M.DData . snd) . Map.toList) (createNewData dataInput Map.empty)
   where
     (binds, dataInput) = splitBindsAndData o
 
 -- | Splits the output into binds and data declaration components (used in createNewData)
 splitBindsAndData :: Output -> ([M.Bind], [(Ident, M.Type, T.Data)])
-splitBindsAndData output = foldl
+splitBindsAndData output =
+    foldl
         ( \(oBinds, oData) (ident, o) -> case o of
             Marked -> error "internal bug in monomorphizer"
             Complete b -> (b : oBinds, oData)
-                                    Data     t d -> (oBinds, (ident, t, d):oData))
+            Data t d -> (oBinds, (ident, t, d) : oData)
+        )
         ([], [])
         (Map.toList output)
 
@@ -325,10 +377,17 @@ createNewData :: [(Ident, M.Type, T.Data)] -> Map.Map Ident M.Data -> Map.Map Id
 createNewData [] o = o
 createNewData ((consIdent, consType, polyData) : input) o =
     createNewData input $
-    Map.insertWith (\_ (M.Data _ cs) -> M.Data newDataType (newCons:cs))
+        Map.insertWith
-      newDataName (M.Data newDataType [newCons]) o
+            (\_ (M.Data _ cs) -> M.Data newDataType (newCons : cs))
+            newDataName
+            (M.Data newDataType [newCons])
+            o
   where
-  T.Data (T.TData polyDataIdent _) _ = polyData
+    polyDataIdent = case polyData of
+        T.Data (T.TData i _) _ -> i
+        T.Data (T.TLit i) _ -> i
+        t -> error $ "Data type is :" ++ show t ++ " which should be impossible"
+
     newDataType = getDataType consType
     newDataName = newName newDataType polyDataIdent
     newCons = M.Inj consIdent consType
@@ -338,4 +397,3 @@ getDataType :: M.Type -> M.Type
 getDataType (M.TFun t1 t2) = getDataType t2
 getDataType tData@(M.TData _ _) = tData
 getDataType _ = error "???"
-