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,36 +39,47 @@ 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
 
 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
 -- can treat them specially.
 data Tag
   = BuiltinTag BuiltinDataTag
   | UserTag TagUser
-  deriving stock (Eq, Generic, Ord, Show)
+  deriving stock (Eq, Generic, Ord)
 
 instance Hashable Tag
 
 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,14 +60,26 @@ data BuiltinDataTag
   | TagBind
   | TagWrite
   | TagReadLn
-  deriving stock (Eq, Generic, Ord, Show)
+  deriving stock (Eq, Generic, Ord)
 
 instance Hashable BuiltinDataTag
 
 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