wip

ollef · Sep 3, 2024 · 47d0a4d · 47d0a4d
1 parent 3398131
commit 47d0a4d
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 38 deletions.
diff --git a/src/Low/Pretty.hs b/src/Low/Pretty.hs
@@ -78,23 +78,35 @@ emptyM module_ = do
 prettyTerm :: Environment v -> Syntax.Term v -> Doc ann
 prettyTerm env = \case
   Syntax.Operand operand -> prettyOperand env operand
-  term@Syntax.Let {} ->
-    line <> indent 2 (prettySeq env term)
-  term@Syntax.Seq {} ->
-    line <> indent 2 (prettySeq env term)
+  Syntax.Let passBy name operation body -> do
+    let (env', name') = extend env name
+    "let"
+      <+> pretty passBy
+      <+> pretty name'
+      <+> "="
+      <+> prettyLetOperation env operation
+        <> line
+        <> prettyTerm env' body
+  Syntax.Seq operation body ->
+    prettySeqOperation env operation
+      <> line
+      <> prettyTerm env body
+
+prettyLetOperation :: Environment v -> Syntax.LetOperation v -> Doc ann
+prettyLetOperation env = \case
   Syntax.Case scrutinee branches defaultBranch ->
     "case"
       <+> prettyOperand env scrutinee
       <+> "of"
-      <> line
-      <> indent
-        2
-        ( vcat $
-            (prettyBranch env <$> branches)
-              <> [ "_" <+> "->" <+> prettyTerm env branch
-                 | Just branch <- [defaultBranch]
-                 ]
-        )
+        <> line
+        <> indent
+          2
+          ( vcat $
+              (prettyBranch env <$> branches)
+                <> [ "_" <+> "->" <+> prettyTerm env branch
+                   | Just branch <- [defaultBranch]
+                   ]
+          )
   Syntax.Call function args ->
     "call"
       <+> prettyLoweredGlobal env function
@@ -109,34 +121,20 @@ prettyTerm env = \case
     "pointer_tag" <+> prettyOperand env operand
   Syntax.Offset operand1 operand2 ->
     prettyOperand env operand1 <+> "+" <+> prettyOperand env operand2
-  Syntax.Copy dst src size ->
-    "copy" <+> commaSep [prettyOperand env dst, prettyOperand env src, prettyOperand env size]
-  Syntax.Store dst src repr ->
-    "store" <+> commaSep [prettyOperand env dst, pretty repr <+> prettyOperand env src]
   Syntax.Load src repr ->
     "load" <+> pretty repr <+> prettyOperand env src
+
+prettySeqOperation :: Environment v -> Syntax.SeqOperation v -> Doc ann
+prettySeqOperation env = \case
+  Syntax.Store dst src repr ->
+    "store" <+> commaSep [prettyOperand env dst, pretty repr <+> prettyOperand env src]
+  Syntax.Copy dst src size ->
+    "copy" <+> commaSep [prettyOperand env dst, prettyOperand env src, prettyOperand env size]
   Syntax.IncreaseReferenceCount operand repr ->
     "increase_reference_count" <+> pretty repr <+> prettyOperand env operand
   Syntax.DecreaseReferenceCount operand repr ->
     "decrease_reference_count" <+> pretty repr <+> prettyOperand env operand
 
-prettySeq :: Environment v -> Syntax.Term v -> Doc ann
-prettySeq env = \case
-  Syntax.Let passBy name term body -> do
-    let (env', name') = extend env name
-    "let"
-      <+> pretty passBy
-      <+> pretty name'
-      <+> "="
-      <+> prettyTerm env term
-      <> line
-      <> prettySeq env' body
-  Syntax.Seq term1 term2 ->
-    prettyTerm env term1
-      <> line
-      <> prettySeq env term2
-  term -> prettyTerm env term
-
 commaSep :: [Doc ann] -> Doc ann
 commaSep = hcat . punctuate (comma <> space)
 
@@ -215,11 +213,11 @@ prettyDefinition env name = \case
 
 prettyFunction :: Environment v -> Syntax.Function v -> Doc ann
 prettyFunction env = \case
-  Syntax.Body passReturnBy body -> "." <+> pretty passReturnBy <+> prettyTerm env body
+  Syntax.Body passReturnBy body -> "." <+> pretty passReturnBy <> line <> indent 2 (prettyTerm env body)
   Syntax.Parameter name passArgBy function' -> do
     let (env', name') = extend env name
     "("
       <> pretty passArgBy
       <+> pretty name'
-      <> ")"
-      <> prettyFunction env' function'
+        <> ")"
+        <> prettyFunction env' function'
diff --git a/src/ReferenceCounting.hs b/src/ReferenceCounting.hs
@@ -232,7 +232,7 @@ referenceCount passBy value = case value of
     Tag _ -> pure (value, Nothing)
     Undefined _ -> pure (value, Nothing)
   Let passValBy name var dead operation body -> do
-    (operation', maybeOperationProvenance, decreaseAfters) <- referenceCountLetOperation passValBy val
+    (operation', maybeOperationProvenance, decreaseAfters) <- referenceCountLetOperation passValBy operation
     forM_ maybeOperationProvenance \valProvenance ->
       modify \s -> s {provenances = EnumMap.insert var valProvenance s.provenances}
     decreaseVar <- case dead of