fix printing & local variables

anoma · May 27, 2024 · 9a5a0e0 · 9a5a0e0
1 parent 20715ab
commit 9a5a0e0
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 11 deletions.
diff --git a/src/Juvix/Compiler/Backend/Rust/Language.hs b/src/Juvix/Compiler/Backend/Rust/Language.hs
@@ -48,6 +48,7 @@ data Let = Let
 
 data ConstDecl = ConstDecl
   { _constVariable :: Text,
+    _constType :: Type,
     _constValue :: Expression
   }
 

diff --git a/src/Juvix/Compiler/Backend/Rust/Pretty/Base.hs b/src/Juvix/Compiler/Backend/Rust/Pretty/Base.hs
@@ -38,7 +38,7 @@ ppBlock :: (Member (Reader Options) r) => [Statement] -> Sem r (Doc Ann)
 ppBlock stmts = do
   stmts' <- mapM ppCode stmts
   let stmts'' = punctuate semi stmts'
-  return $ braces (indent' (vsep stmts''))
+  return $ oneLineOrNextBraces (vsep stmts'')
 
 instance PrettyCode Type where
   ppCode = \case
@@ -51,17 +51,17 @@ instance PrettyCode FunctionArgument where
     ty <- ppCode _functionArgumentType
     n <- ppName KNameLocal _functionArgumentName
     mut <- ppMut _functionArgumentMutable
-    return $ mut <?+> n <> ":" <+> ty
+    return $ mut <?+> n <> colon <+> ty
 
 instance PrettyCode Function where
   ppCode Function {..} = do
     attrs <- ppAttrs _functionAttributes
     name <- ppName KNameFunction _functionName
     args <- mapM ppCode _functionArguments
-    rty <- maybe (return Nothing) (ppCode >=> return . Just) _functionReturnType
+    rty <- maybe (return Nothing) (ppCode >=> return . Just . ("->" <+>)) _functionReturnType
     body <- ppBlock _functionBody
-    let args' = punctuate colon args
-    return $ attrs <> kwFn <+> name <> parens (hsep args') <+> rty <?+> "=" <+> body
+    let args' = punctuate comma args
+    return $ attrs <> kwFn <+> name <> parens (hsep args') <+> rty <?+> body <> line
 
 instance PrettyCode Statement where
   ppCode = \case
@@ -86,8 +86,9 @@ instance PrettyCode Let where
 instance PrettyCode ConstDecl where
   ppCode ConstDecl {..} = do
     name <- ppName KNameLocal _constVariable
+    ty <- ppCode _constType
     val <- ppCode _constValue
-    return $ kwConst <+> name <+> "=" <+> val
+    return $ kwConst <+> name <> colon <+> ty <+> "=" <+> val
 
 instance PrettyCode Assignment where
   ppCode Assignment {..} = do
@@ -114,11 +115,11 @@ instance PrettyCode Match where
   ppCode Match {..} = do
     val <- ppCode _matchValue
     brs <- mapM ppCode _matchBranches
-    return $ kwMatch <+> val <+> braces (indent' (vsep brs))
+    return $ kwMatch <+> val <+> oneLineOrNextBraces (vsep brs)
 
 instance PrettyCode Loop where
   ppCode Loop {..} = do
-    let lab = fmap pretty _loopLabel
+    let lab = fmap ((<> colon) . pretty) _loopLabel
     body <- ppBlock _loopBody
     return $ lab <?+> kwLoop <+> body
 
@@ -176,7 +177,7 @@ instance PrettyCode Program where
         \#[allow(unused_imports)]\n\
         \use juvix::defs::*;\n\
         \#[allow(unused_imports)]\n\
-        \use juvix::apply;\
+        \use juvix::apply;\n\
         \#[allow(unused_imports)]\n\
         \use juvix::tapply;\n\
         \#[allow(unused_imports)]\n\

diff --git a/src/Juvix/Compiler/Backend/Rust/Translation/FromReg.hs b/src/Juvix/Compiler/Backend/Rust/Translation/FromReg.hs
@@ -31,7 +31,7 @@ fromReg lims tab =
           _functionArguments = [],
           _functionReturnType = Nothing,
           _functionBody =
-            [ stmtLet NotMut "result" (mkCall "program" [ExprVerbatim "&mut Memory::new()", mkVec []]),
+            [ stmtLet NotMut "result" (mkCall "program" [ExprVerbatim "&mut Memory::new()", mkInteger @Int 0, mkVec []]),
               StatementExpression (mkCall "println!" [mkString "{}", mkVar "result"]),
               StatementReturn (Return Nothing)
             ]
@@ -62,6 +62,7 @@ fromReg lims tab =
           StatementConst $
             ConstDecl
               { _constVariable = getFunctionIdent info (funInfo ^. Reg.functionSymbol),
+                _constType = Word,
                 _constValue = mkInteger (getFUID info (funInfo ^. Reg.functionSymbol))
               }
 
@@ -93,8 +94,14 @@ fromRegFunction :: Reg.ExtraInfo -> Reg.FunctionInfo -> MatchBranch
 fromRegFunction info funInfo =
   MatchBranch
     { _matchBranchPattern = Just $ mkVar $ getFunctionIdent info (funInfo ^. Reg.functionSymbol),
-      _matchBranchBody = fromRegCode info (funInfo ^. Reg.functionCode)
+      _matchBranchBody = varDecls ++ fromRegCode info (funInfo ^. Reg.functionCode)
     }
+  where
+    varsNum :: Int
+    varsNum = getLocalVarsNum info (funInfo ^. Reg.functionSymbol)
+
+    varDecls :: [Statement]
+    varDecls = map (\n -> stmtDecl Mut ("var_" <> show n)) [0 .. varsNum - 1]
 
 fromRegCode :: Reg.ExtraInfo -> Reg.Code -> [Statement]
 fromRegCode info = concatMap (fromRegInstr info)

diff --git a/src/Juvix/Compiler/Backend/Rust/Translation/FromReg/Base.hs b/src/Juvix/Compiler/Backend/Rust/Translation/FromReg/Base.hs
@@ -64,6 +64,15 @@ stmtLet isMut result value =
         _letInitializer = Just value
       }
 
+stmtDecl :: IsMut -> Text -> Statement
+stmtDecl isMut result =
+  StatementLet $
+    Let
+      { _letMutable = isMut,
+        _letVariable = result,
+        _letInitializer = Nothing
+      }
+
 stmtIf :: Expression -> [Statement] -> [Statement] -> Statement
 stmtIf v br1 br2 = StatementIf $ If v br1 br2