Fix JuvixTree parsing and pretty printing (#3024)

Recent changes to the compiler left JuvixTree parsing and pretty
printing not in sync with each other.

src/Juvix/Compiler/Core/Data/InfoTable.hs

@@ -106,11 +106,13 @@ typeName' :: InfoTable -> Symbol -> Text
 typeName' tab sym = lookupTabInductiveInfo tab sym ^. inductiveName
 
 identNames' :: InfoTable -> HashSet Text
-identNames' tab =
-  HashSet.fromList $
-    map (^. identifierName) (HashMap.elems (tab ^. infoIdentifiers))
-      ++ map (^. constructorName) (HashMap.elems (tab ^. infoConstructors))
-      ++ map (^. inductiveName) (HashMap.elems (tab ^. infoInductives))
+identNames' = HashSet.fromList . identNamesList'
+
+identNamesList' :: InfoTable -> [Text]
+identNamesList' tab =
+  map (^. identifierName) (HashMap.elems (tab ^. infoIdentifiers))
+    ++ map (^. constructorName) (HashMap.elems (tab ^. infoConstructors))
+    ++ map (^. inductiveName) (HashMap.elems (tab ^. infoInductives))
 
 freshIdentName' :: InfoTable -> Text -> Text
 freshIdentName' tab = freshName (identNames' tab)

src/Juvix/Compiler/Core/Data/Module.hs

@@ -110,6 +110,9 @@ constrName md tag = lookupConstructorInfo md tag ^. constructorName
 identNames :: Module -> HashSet Text
 identNames m = identNames' (computeCombinedInfoTable m)
 
+identNamesList :: Module -> [Text]
+identNamesList m = identNamesList' (computeCombinedInfoTable m)
+
 freshIdentName :: Module -> Text -> Text
 freshIdentName m = freshName (identNames m)
 

src/Juvix/Compiler/Core/Data/TransformationId.hs

@@ -75,7 +75,7 @@ toVampIRTransformations =
 
 toStrippedTransformations :: TransformationId -> [TransformationId]
 toStrippedTransformations checkId =
-  combineInfoTablesTransformations ++ [checkId, LambdaLetRecLifting, TopEtaExpand, OptPhaseExec, MoveApps, RemoveTypeArgs]
+  combineInfoTablesTransformations ++ [checkId, LambdaLetRecLifting, TopEtaExpand, OptPhaseExec, MoveApps, RemoveTypeArgs, DisambiguateNames]
 
 instance TransformationId' TransformationId where
   transformationText :: TransformationId -> Text

src/Juvix/Compiler/Core/Language/Base.hs

@@ -39,14 +39,11 @@ instance Show Symbol where
 defaultSymbol :: Word -> Symbol
 defaultSymbol = Symbol defaultModuleId
 
-uniqueName :: Text -> Symbol -> Text
-uniqueName txt sym = txt <> "_" <> show sym
-
 data TagUser = TagUser
   { _tagUserModuleId :: ModuleId,
     _tagUserWord :: Word
   }
-  deriving stock (Eq, Generic, Ord, Show)
+  deriving stock (Eq, Generic, Ord)
 
 instance Hashable TagUser
 
@@ -54,6 +51,12 @@ instance Serialize TagUser
 
 instance NFData TagUser
 
+instance Pretty TagUser where
+  pretty TagUser {..} = pretty _tagUserWord <> "@" <> pretty _tagUserModuleId
+
+instance Show TagUser where
+  show = show . pretty
+
 -- | Tag of a constructor, uniquely identifying it. Tag values are consecutive
 -- and separate from symbol IDs. We might need fixed special tags in Core for
 -- common "builtin" constructors, e.g., unit, nat, so that the code generator
@@ -61,7 +64,7 @@ instance NFData TagUser
 data Tag
   = BuiltinTag BuiltinDataTag
   | UserTag TagUser
-  deriving stock (Eq, Generic, Ord, Show)
+  deriving stock (Eq, Generic, Ord)
 
 instance Hashable Tag
 
@@ -69,6 +72,14 @@ instance Serialize Tag
 
 instance NFData Tag
 
+instance Pretty Tag where
+  pretty = \case
+    BuiltinTag b -> pretty b
+    UserTag u -> pretty u
+
+instance Show Tag where
+  show = show . pretty
+
 isBuiltinTag :: Tag -> Bool
 isBuiltinTag = \case
   BuiltinTag {} -> True

src/Juvix/Compiler/Core/Language/Builtins.hs

@@ -1,7 +1,10 @@
 module Juvix.Compiler.Core.Language.Builtins where
 
+import GHC.Show qualified as Show
 import Juvix.Extra.Serialize
+import Juvix.Extra.Strings qualified as Str
 import Juvix.Prelude
+import Prettyprinter
 
 -- Builtin operations which the evaluator and the code generator treat
 -- specially and non-uniformly.
@@ -57,7 +60,7 @@ data BuiltinDataTag
   | TagBind
   | TagWrite
   | TagReadLn
-  deriving stock (Eq, Generic, Ord, Show)
+  deriving stock (Eq, Generic, Ord)
 
 instance Hashable BuiltinDataTag
 
@@ -65,6 +68,18 @@ instance Serialize BuiltinDataTag
 
 instance NFData BuiltinDataTag
 
+instance Pretty BuiltinDataTag where
+  pretty = \case
+    TagTrue -> Str.true_
+    TagFalse -> Str.false_
+    TagReturn -> Str.return
+    TagBind -> Str.bind
+    TagWrite -> Str.write
+    TagReadLn -> Str.readLn
+
+instance Show BuiltinDataTag where
+  show = show . pretty
+
 builtinOpArgsNum :: BuiltinOp -> Int
 builtinOpArgsNum = \case
   OpIntAdd -> 2

src/Juvix/Compiler/Core/Transformation/DisambiguateNames.hs

@@ -122,6 +122,28 @@ disambiguateNodeNames md = disambiguateNodeNames' disambiguate md
     names :: HashSet Text
     names = identNames md
 
+disambiguateTopNames :: Module -> Module
+disambiguateTopNames md =
+  mapInductives (\i -> over inductiveName (renameDuplicated (i ^. inductiveSymbol)) i)
+    . mapConstructors (\i -> over constructorName (renameDuplicated (i ^. constructorTag)) i)
+    . mapIdents (\i -> over identifierName (renameDuplicated (i ^. identifierSymbol)) i)
+    $ md
+  where
+    duplicatedNames :: HashSet Text
+    duplicatedNames =
+      HashSet.fromList
+        . map head
+        . filter (\x -> length x > 1)
+        . NonEmpty.group
+        . sort
+        . identNamesList
+        $ md
+
+    renameDuplicated :: (Show a) => a -> Text -> Text
+    renameDuplicated sym name
+      | HashSet.member name duplicatedNames = uniqueName name sym
+      | otherwise = name
+
 setArgNames :: Module -> Symbol -> Node -> Node
 setArgNames md sym node = reLambdas lhs' body
   where
@@ -135,8 +157,9 @@ setArgNames md sym node = reLambdas lhs' body
 
 disambiguateNames :: Module -> Module
 disambiguateNames md =
-  let md' = mapT (setArgNames md) md
-   in mapAllNodes (disambiguateNodeNames md') md'
+  let md1 = disambiguateTopNames md
+      md2 = mapT (setArgNames md1) md1
+   in mapAllNodes (disambiguateNodeNames md2) md2
 
 disambiguateNames' :: InfoTable -> InfoTable
 disambiguateNames' = withInfoTable disambiguateNames

src/Juvix/Compiler/Core/Translation/Stripped/FromCore.hs

@@ -136,10 +136,7 @@ translateFunctionInfo tab IdentifierInfo {..} =
       { _functionName = _identifierName,
         _functionLocation = _identifierLocation,
         _functionSymbol = _identifierSymbol,
-        _functionBody =
-          translateFunction
-            _identifierArgsNum
-            body,
+        _functionBody = translateFunction _identifierArgsNum body,
         _functionType = translateType _identifierType,
         _functionArgsNum = _identifierArgsNum,
         _functionArgsInfo = map translateArgInfo (lambdaBinders body),

src/Juvix/Compiler/Tree/Keywords.hs

@@ -9,6 +9,8 @@ import Juvix.Compiler.Tree.Keywords.Base
 import Juvix.Data.Keyword.All
   ( kwAdd_,
     kwAlloc,
+    kwAnomaByteArrayFromAnomaContents,
+    kwAnomaByteArrayToAnomaContents,
     kwAnomaDecode,
     kwAnomaEncode,
     kwAnomaGet,
@@ -35,6 +37,9 @@ import Juvix.Data.Keyword.All
     kwFieldDiv,
     kwFieldMul,
     kwFieldSub,
+    kwFieldToInt,
+    kwIntToField,
+    kwIntToUInt8,
     kwLe_,
     kwLt_,
     kwMod_,
@@ -47,6 +52,7 @@ import Juvix.Data.Keyword.All
     kwStrcat,
     kwSub_,
     kwTrace,
+    kwUInt8ToInt,
   )
 import Juvix.Prelude
 
@@ -90,9 +96,15 @@ allKeywords =
          kwAnomaSign,
          kwAnomaSignDetached,
          kwAnomaVerifyWithMessage,
+         kwAnomaByteArrayFromAnomaContents,
+         kwAnomaByteArrayToAnomaContents,
          kwPoseidon,
          kwEcOp,
          kwRandomEcPoint,
          kwByteArrayLength,
-         kwByteArrayFromListUInt8
+         kwByteArrayFromListUInt8,
+         kwIntToUInt8,
+         kwUInt8ToInt,
+         kwIntToField,
+         kwFieldToInt
        ]

src/Juvix/Compiler/Tree/Pretty/Base.hs

@@ -107,20 +107,13 @@ instance PrettyCode Value where
 
 instance PrettyCode TypeInductive where
   ppCode :: (Member (Reader Options) r) => TypeInductive -> Sem r (Doc Ann)
-  ppCode TypeInductive {..} = do
-    names <- asks (^. optSymbolNames)
-    let name = fromJust (HashMap.lookup _typeInductiveSymbol names)
-    return $ annotate (AnnKind KNameInductive) (pretty name)
+  ppCode TypeInductive {..} = ppIndName _typeInductiveSymbol
 
 instance PrettyCode TypeConstr where
   ppCode :: (Member (Reader Options) r) => TypeConstr -> Sem r (Doc Ann)
   ppCode TypeConstr {..} = do
-    symNames <- asks (^. optSymbolNames)
-    let indname = fromJust (HashMap.lookup _typeConstrInductive symNames)
-    let iname = annotate (AnnKind KNameInductive) (pretty indname)
-    tagNames <- asks (^. optTagNames)
-    let ctrname = fromJust (HashMap.lookup _typeConstrTag tagNames)
-    let cname = annotate (AnnKind KNameConstructor) (pretty ctrname)
+    iname <- ppIndName _typeConstrInductive
+    cname <- ppConstrName _typeConstrTag
     args <- mapM ppCode _typeConstrFields
     return $ iname <> kwColon <> cname <> parens (hsep (punctuate comma args))
 
@@ -442,7 +435,7 @@ instance (PrettyCode a) => PrettyCode [a] where
 ppFunInfo :: (Member (Reader Options) r) => (t -> Sem r (Doc Ann)) -> FunctionInfo' t e -> Sem r (Doc Ann)
 ppFunInfo ppCode' FunctionInfo {..} = do
   argtys <- mapM ppCode (take _functionArgsNum (typeArgs _functionType))
-  let argnames = map (fmap variable) _functionArgNames
+  let argnames = map (fmap (variable . quoteName)) _functionArgNames
       args = zipWithExact (\mn ty -> maybe mempty (\n -> n <+> colon <> space) mn <> ty) argnames argtys
   targetty <- ppCode (if _functionArgsNum == 0 then _functionType else typeTarget _functionType)
   c <- ppCode' _functionCode

src/Juvix/Compiler/Tree/Pretty/Extra.hs

@@ -1,6 +1,5 @@
 module Juvix.Compiler.Tree.Pretty.Extra where
 
-import Data.Text qualified as Text
 import Juvix.Data.CodeAnn
 import Juvix.Prelude
 
@@ -17,21 +16,33 @@ variable :: Text -> Doc Ann
 variable a = annotate (AnnKind KNameLocal) (pretty a)
 
 quoteName :: Text -> Text
-quoteName txt =
-  foldr
-    (uncurry Text.replace)
-    txt
-    [ ("$", "__dollar__"),
-      (":", "__colon__"),
-      ("@", "__at__"),
-      ("arg", "__arg__"),
-      ("tmp", "__tmp__")
-    ]
+quoteName =
+  quote1 . quote0
+  where
+    quote0 :: Text -> Text
+    quote0 =
+      replaceSubtext
+        [ ("$", "__dollar__"),
+          (":", "__colon__"),
+          ("@", "__at__"),
+          (".", "__dot__"),
+          (",", "__comma__"),
+          (";", "__semicolon__")
+        ]
+
+    quote1 :: Text -> Text
+    quote1 =
+      replaceText
+        [ ("arg", "__arg__"),
+          ("tmp", "__tmp__"),
+          ("sub", "__sub__"),
+          ("add", "__add__"),
+          ("mul", "__mul__"),
+          ("div", "__div__")
+        ]
 
 quoteFunName :: Text -> Text
-quoteFunName txt =
-  foldr
-    (uncurry Text.replace)
-    txt
+quoteFunName =
+  replaceText
     [ ("readLn", "__readLn__")
     ]

src/Juvix/Compiler/Tree/Translation/FromSource.hs

@@ -110,6 +110,10 @@ parseUnop =
     <|> parseUnaryOp kwTrace OpTrace
     <|> parseUnaryOp kwFail OpFail
     <|> parseUnaryOp kwArgsNum (PrimUnop OpArgsNum)
+    <|> parseUnaryOp kwIntToUInt8 (PrimUnop OpIntToUInt8)
+    <|> parseUnaryOp kwUInt8ToInt (PrimUnop OpUInt8ToInt)
+    <|> parseUnaryOp kwIntToField (PrimUnop OpIntToField)
+    <|> parseUnaryOp kwFieldToInt (PrimUnop OpFieldToInt)
 
 parseUnaryOp ::
   (Members '[Reader ParserSig, InfoTableBuilder, State LocalParams] r) =>
@@ -149,6 +153,8 @@ parseAnoma =
     <|> parseAnoma' kwAnomaSign OpAnomaSign
     <|> parseAnoma' kwAnomaSignDetached OpAnomaSignDetached
     <|> parseAnoma' kwAnomaVerifyWithMessage OpAnomaVerifyWithMessage
+    <|> parseAnoma' kwAnomaByteArrayToAnomaContents OpAnomaByteArrayToAnomaContents
+    <|> parseAnoma' kwAnomaByteArrayFromAnomaContents OpAnomaByteArrayFromAnomaContents
 
 parseAnoma' ::
   (Members '[Reader ParserSig, InfoTableBuilder, State LocalParams] r) =>

src/Juvix/Compiler/Tree/Translation/FromSource/Base.hs

@@ -91,10 +91,10 @@ declareBuiltins = do
   sym <- lift $ freshSymbol' @t @e
   let tyio = mkTypeInductive sym
       constrs =
-        [ createBuiltinConstr sym TagReturn "return" (mkTypeFun [TyDynamic] tyio) i,
-          createBuiltinConstr sym TagBind "bind" (mkTypeFun [tyio, mkTypeFun [TyDynamic] tyio] tyio) i,
-          createBuiltinConstr sym TagWrite "write" (mkTypeFun [TyDynamic] tyio) i,
-          createBuiltinConstr sym TagReadLn "readLn" tyio i
+        [ createBuiltinConstr sym TagReturn (show TagReturn) (mkTypeFun [TyDynamic] tyio) i,
+          createBuiltinConstr sym TagBind (show TagBind) (mkTypeFun [tyio, mkTypeFun [TyDynamic] tyio] tyio) i,
+          createBuiltinConstr sym TagWrite (show TagWrite) (mkTypeFun [TyDynamic] tyio) i,
+          createBuiltinConstr sym TagReadLn (show TagReadLn) tyio i
         ]
   lift $
     registerInductive' @t @e
@@ -310,6 +310,7 @@ typeNamed = do
     "string" -> return TyString
     "unit" -> return TyUnit
     "uint8" -> return mkTypeUInt8
+    "bytearray" -> return TyByteArray
     _ -> do
       idt <- lift $ getIdent' @t @e txt
       case idt of

src/Juvix/Data/Keyword/All.hs

@@ -286,6 +286,18 @@ kwPrealloc = asciiKw Str.prealloc
 kwArgsNum :: Keyword
 kwArgsNum = asciiKw Str.instrArgsNum
 
+kwIntToUInt8 :: Keyword
+kwIntToUInt8 = asciiKw Str.instrIntToUInt8
+
+kwUInt8ToInt :: Keyword
+kwUInt8ToInt = asciiKw Str.instrUInt8ToInt
+
+kwIntToField :: Keyword
+kwIntToField = asciiKw Str.instrIntToField
+
+kwFieldToInt :: Keyword
+kwFieldToInt = asciiKw Str.instrFieldToInt
+
 kwByteArrayFromListUInt8 :: Keyword
 kwByteArrayFromListUInt8 = asciiKw Str.instrByteArrayFromListUInt8
 
@@ -478,6 +490,12 @@ kwAnomaVerifyWithMessage = asciiKw Str.anomaVerifyWithMessage
 kwByteArrayFromListByte :: Keyword
 kwByteArrayFromListByte = asciiKw Str.byteArrayFromListByte
 
+kwAnomaByteArrayToAnomaContents :: Keyword
+kwAnomaByteArrayToAnomaContents = asciiKw Str.anomaByteArrayToAnomaContents
+
+kwAnomaByteArrayFromAnomaContents :: Keyword
+kwAnomaByteArrayFromAnomaContents = asciiKw Str.anomaByteArrayFromAnomaContents
+
 kwByteArrayLength :: Keyword
 kwByteArrayLength = asciiKw Str.byteArrayLength
 

src/Juvix/Extra/Strings.hs

@@ -614,6 +614,18 @@ false_ = "false"
 default_ :: (IsString s) => s
 default_ = "default"
 
+return :: (IsString s) => s
+return = "return"
+
+bind :: (IsString s) => s
+bind = "bind"
+
+write :: (IsString s) => s
+write = "write"
+
+readLn :: (IsString s) => s
+readLn = "readLn"
+
 plus :: (IsString s) => s
 plus = "+"
 

src/Juvix/Prelude/Base/Foundation.hs

@@ -310,6 +310,15 @@ isFirstLetter = \case
   h : _ -> isLetter h
   _ -> False
 
+uniqueName :: (Show a) => Text -> a -> Text
+uniqueName txt sym = txt <> "_" <> show sym
+
+replaceSubtext :: [(Text, Text)] -> Text -> Text
+replaceSubtext texts txt = foldr (uncurry Text.replace) txt texts
+
+replaceText :: [(Text, Text)] -> Text -> Text
+replaceText texts txt = fromMaybe txt (HashMap.lookup txt (HashMap.fromList texts))
+
 --------------------------------------------------------------------------------
 -- Foldable
 --------------------------------------------------------------------------------

test/Tree.hs

@@ -3,7 +3,8 @@ module Tree where
 import Base
 import Tree.Asm qualified as Asm
 import Tree.Eval qualified as Eval
+import Tree.Parse qualified as Parse
 import Tree.Transformation qualified as Transformation
 
 allTests :: TestTree
-allTests = testGroup "JuvixTree tests" [Eval.allTests, Asm.allTests, Transformation.allTests]
+allTests = testGroup "JuvixTree tests" [Parse.allTests, Eval.allTests, Asm.allTests, Transformation.allTests]

test/Tree/Parse.hs

@@ -0,0 +1,7 @@
+module Tree.Parse where
+
+import Base
+import Tree.Parse.Positive qualified as P
+
+allTests :: TestTree
+allTests = testGroup "JuvixTree parsing" [P.allTests]

test/Tree/Parse/Base.hs

@@ -0,0 +1,32 @@
+module Tree.Parse.Base where
+
+import Base
+import Juvix.Compiler.Tree.Data.InfoTable
+import Juvix.Compiler.Tree.Pretty
+import Juvix.Compiler.Tree.Translation.FromSource
+import Juvix.Data.PPOutput
+
+treeParseAssertion :: Path Abs File -> (String -> IO ()) -> Assertion
+treeParseAssertion mainFile step = do
+  step "Parse"
+  r <- parseFile mainFile
+  case r of
+    Left err -> assertFailure (prettyString err)
+    Right tab -> do
+      withTempDir'
+        ( \dirPath -> do
+            let outputFile = dirPath <//> $(mkRelFile "out.out")
+            step "Print"
+            writeFileEnsureLn outputFile (ppPrint tab tab)
+            step "Parse printed"
+            r' <- parseFile outputFile
+            case r' of
+              Left err -> assertFailure (prettyString err)
+              Right tab' -> do
+                assertBool ("Check: print . parse = print . parse . print . parse") (ppPrint tab tab == ppPrint tab' tab')
+        )
+
+parseFile :: Path Abs File -> IO (Either MegaparsecError InfoTable)
+parseFile f = do
+  s <- readFile f
+  return (runParser f s)

test/Tree/Parse/Positive.hs

@@ -0,0 +1,23 @@
+module Tree.Parse.Positive where
+
+import Base
+import Tree.Eval.Positive qualified as Eval
+import Tree.Parse.Base
+
+type PosTest = Eval.PosTest
+
+testDescr :: PosTest -> TestDescr
+testDescr Eval.PosTest {..} =
+  let tRoot = Eval.root <//> _relDir
+      file' = tRoot <//> _file
+   in TestDescr
+        { _testName = _name,
+          _testRoot = tRoot,
+          _testAssertion = Steps $ treeParseAssertion file'
+        }
+
+allTests :: TestTree
+allTests =
+  testGroup
+    "JuvixTree parsing positive tests"
+    (map (mkTest . testDescr) Eval.tests)