Store arg passing info in low functions

src/Compiler.hs

@@ -35,9 +35,8 @@ compile assemblyDir saveAssembly outputExecutableFile maybeOptimisationLevel pri
       emptyPrettyEnv <- Pretty.emptyM moduleName
       forM_ defNames \defName -> do
         maybeLoweredDef <- fetch $ Query.LoweredDefinition defName
-        forM_ maybeLoweredDef \loweredDef -> do
-          prettyDef <- Pretty.prettyDefinition emptyPrettyEnv defName loweredDef
-          liftIO $ putDocW 120 $ prettyDef <> line
+        forM_ maybeLoweredDef \loweredDef ->
+          liftIO $ putDocW 120 $ Pretty.prettyDefinition emptyPrettyEnv defName loweredDef <> line
 
     llvmModule <- fetch $ Query.LLVMModule moduleName
     let llvmFileName = moduleAssemblyDir </> toS moduleNameText <.> "ll"

src/Low/Pretty.hs

@@ -14,7 +14,6 @@ import qualified Data.Kind
 import qualified Data.Sequence as Seq
 import qualified Data.Text.Unsafe as Text
 import Index
-import Low.PassBy (PassBy)
 import qualified Low.Syntax as Syntax
 import Name (Name (Name))
 import qualified Name
@@ -196,26 +195,20 @@ prettyBranch env = \case
 
 -------------------------------------------------------------------------------
 
-prettyDefinition :: (MonadFetch Query m) => Environment Void -> Name.Lifted -> Syntax.Definition -> m (Doc ann)
-prettyDefinition env name def = do
-  signature <- fetch $ Query.LowSignature name
-  pure case (def, signature) of
-    (Syntax.ConstantDefinition term, Syntax.ConstantSignature repr) ->
-      prettyLiftedGlobal env name <+> pretty repr <+> "=" <+> prettyTerm env term
-    (Syntax.ConstantDefinition _, _) -> panic "definition signature mismatch"
-    (Syntax.FunctionDefinition function, Syntax.FunctionSignature passArgsBy passReturnBy) ->
-      prettyLiftedGlobal env name <+> "=" <+> "\\" <> prettyFunction env passArgsBy passReturnBy function
-    (Syntax.FunctionDefinition _, _) -> panic "definition signature mismatch"
-
-prettyFunction :: Environment v -> [PassBy] -> PassBy -> Syntax.Function v -> Doc ann
-prettyFunction env passArgsBy passReturnBy function = case (passArgsBy, function) of
-  ([], Syntax.Body body) -> " ->" <+> pretty passReturnBy <+> prettyTerm env body
-  ([], _) -> panic "function signature mismatch"
-  (passArgBy : passArgsBy', Syntax.Parameter name function') -> do
+prettyDefinition :: Environment Void -> Name.Lifted -> Syntax.Definition -> Doc ann
+prettyDefinition env name = \case
+  Syntax.ConstantDefinition repr term ->
+    prettyLiftedGlobal env name <+> pretty repr <+> "=" <+> prettyTerm env term
+  Syntax.FunctionDefinition function ->
+    prettyLiftedGlobal env name <+> "=" <+> "\\" <> prettyFunction env function
+
+prettyFunction :: Environment v -> Syntax.Function v -> Doc ann
+prettyFunction env = \case
+  Syntax.Body passReturnBy body -> " ->" <+> pretty passReturnBy <+> prettyTerm env body
+  Syntax.Parameter name passArgBy function' -> do
     let (env', name') = extend env name
     "("
       <> pretty passArgBy
       <+> pretty name'
         <> ")"
-        <> prettyFunction env' passArgsBy' passReturnBy function'
-  (_ : _, _) -> panic "function signature mismatch"
+        <> prettyFunction env' function'

src/Low/Syntax.hs

@@ -41,14 +41,14 @@ data Branch v
   deriving (Eq, Show, Generic, Hashable)
 
 data Function v
-  = Body !(Term v)
-  | Parameter !Name !(Scope Function v)
+  = Body !PassBy !(Term v)
+  | Parameter !Name !PassBy !(Scope Function v)
   deriving (Eq, Show, Generic, Hashable)
 
 type Type = Term
 
 data Definition
-  = ConstantDefinition !(Term Void)
+  = ConstantDefinition !Representation !(Term Void)
   | FunctionDefinition !(Function Void)
   deriving (Eq, Show, Generic, Hashable)
 

src/Lower.hs

@@ -32,6 +32,7 @@ import qualified Low.Syntax
 import Monad
 import Name (Name)
 import qualified Name
+import Prettyprinter
 import Protolude hiding (nonEmpty, repr)
 import qualified Query
 import Rock.Core
@@ -81,7 +82,7 @@ data Collectible
   | CollectibleSeq !Value
   deriving (Show)
 
-data Function = Function [(Name, Var)] !Value
+data Function = Function [(Name, PassBy, Var)] !PassBy !Value
 
 type Collect = StateT (Tsil Collectible) M
 
@@ -155,18 +156,18 @@ definition name = \case
     constantDefinition term = do
       signature <- fetch $ Query.LowSignature name
       case signature of
-        Low.Syntax.ConstantSignature _ -> do
+        Low.Syntax.ConstantSignature repr -> do
           value <- runCollect $ storeTerm CC.empty mempty (Global name) term
           let term' = readback Index.Map.Empty value
-          pure $ Just $ Low.Syntax.ConstantDefinition term'
+          pure $ Just $ Low.Syntax.ConstantDefinition repr term'
         _ -> panic "Constant without constant signature"
 
     functionDefinition tele = do
       signature <- fetch $ Query.LowSignature name
       case signature of
         Low.Syntax.FunctionSignature passArgsBy passReturnBy -> do
           functionValue <-
-            genRunCollect (\(_, _, result) -> Operand result) (\(params, returns, _) body -> Function (returns <> params) body) $
+            genRunCollect (\(_, _, result) -> Operand result) (\(params, returns, _) body -> Function (returns <> params) passReturnBy body) $
               lowerFunction CC.empty mempty passArgsBy passReturnBy tele
           let function = readbackFunction Index.Map.Empty functionValue
           pure $ Just $ Low.Syntax.FunctionDefinition function
@@ -178,7 +179,7 @@ lowerFunction
   -> [PassBy]
   -> PassBy
   -> Telescope Name CC.Syntax.Type CC.Syntax.Term v
-  -> Collect ([(Name, Var)], [(Name, Var)], Operand)
+  -> Collect ([(Name, PassBy, Var)], [(Name, PassBy, Var)], Operand)
 lowerFunction context indices passArgsBy passReturnBy tele = case (tele, passArgsBy) of
   (Telescope.Empty body, []) -> case passReturnBy of
     PassBy.Value repr -> do
@@ -188,7 +189,7 @@ lowerFunction context indices passArgsBy passReturnBy tele = case (tele, passArg
     PassBy.Reference -> do
       dst <- lift freshVar
       result <- storeTerm context indices (Var dst) body
-      pure ([], [("return", dst)], result)
+      pure ([], [("return", PassBy.Reference, dst)], result)
   (Telescope.Empty _, _) -> panic "Function signature mismatch"
   (Telescope.Extend name type_ _plicity tele', passArgBy : passArgsBy') -> do
     type' <- lift $ Evaluation.evaluate (CC.toEnvironment context) type_
@@ -200,7 +201,7 @@ lowerFunction context indices passArgsBy passReturnBy tele = case (tele, passArg
     (context', var) <- lift $ CC.extend context type'
     let indices' = indices Seq.:|> OperandStorage (Var var) operandRepr
     (params, returns, result) <- lowerFunction context' indices' passArgsBy' passReturnBy tele'
-    pure ((name, var) : params, returns, result)
+    pure ((name, passArgBy, var) : params, returns, result)
   (Telescope.Extend {}, _) -> panic "Function signature mismatch"
 
 storeOperand
@@ -491,7 +492,7 @@ storeCall context indices dst function args passArgsBy passReturnBy = do
         callResult <- letValue repr "call_result" $ Call function callArgs
         storeOperand dst $ OperandStorage callResult $ Value repr
       PassBy.Reference ->
-        letReference "call_result_size" $ Call function (dst : callArgs)
+        letValue Representation.type_ "call_result_size" $ Call function (dst : callArgs)
 
 storeBranch
   :: CC.Context v
@@ -531,13 +532,13 @@ boxedConstructorSize env con params args = do
 -------------------------------------------------------------------------------
 
 readbackFunction :: Index.Map v Var -> Function -> Low.Syntax.Function v
-readbackFunction outerEnv (Function params body) =
+readbackFunction outerEnv (Function params returnRepr body) =
   go outerEnv params
   where
-    go :: Index.Map v Var -> [(Name, Var)] -> Low.Syntax.Function v
-    go env [] = Low.Syntax.Body $ readback env body
-    go env ((name, var) : params') =
-      Low.Syntax.Parameter name $ go (env Index.Map.:> var) params'
+    go :: Index.Map v Var -> [(Name, PassBy, Var)] -> Low.Syntax.Function v
+    go env [] = Low.Syntax.Body returnRepr $ readback env body
+    go env ((name, passParamBy, var) : params') =
+      Low.Syntax.Parameter name passParamBy $ go (env Index.Map.:> var) params'
 
 readback :: Index.Map v Var -> Value -> Low.Syntax.Term v
 readback env = \case