From 510490a5bfe3394ed093b19a162a2937a0a0359c Mon Sep 17 00:00:00 2001 From: Jan Mas Rovira Date: Tue, 23 Jan 2024 14:36:19 +0100 Subject: [PATCH] Support `MemRepTuple` in the Nockma backend (#2586) We can represent Anoma types like [resource](https://github.com/anoma/anoma/blob/e18e50e3c380b2867da6cb81f427334eb5c0b2d9/hoon/resource-machine.hoon#L7) as Juvix records. The Nockma encoding of types uses Nockma 'tuples' where each component of the tuple holds a value of a field. So for Juvix->Anoma integration it is convenient to compile values of record types as Nockma tuples. We already have the concept of representing constructors of inductive types that have only one non-zero-field constructor in the compiler, see [`MemRepTuple`](https://github.com/anoma/juvix/blob/1147e1fce131f5b423fd558e5636e6aaf23120ac/src/Juvix/Compiler/Tree/Language/Rep.hs#L13). In this PR, as part of the Nockma step, we mark constructors that satisfy the requirements of the `MemRepTuple` translation as such. Then we use a tuple encoding for those constructors. --------- Co-authored-by: Paul Cadman --- .../Compiler/Core/Data/Stripped/InfoTable.hs | 2 + .../Core/Translation/Stripped/FromCore.hs | 1 + .../Compiler/Nockma/Translation/FromAsm.hs | 148 ++++++++++++++---- .../Compiler/Tree/Translation/FromCore.hs | 2 +- test/Nockma/Compile/Positive.hs | 74 ++++++--- 5 files changed, 176 insertions(+), 51 deletions(-) diff --git a/src/Juvix/Compiler/Core/Data/Stripped/InfoTable.hs b/src/Juvix/Compiler/Core/Data/Stripped/InfoTable.hs index 14022a8ec9..c3e9e78cac 100644 --- a/src/Juvix/Compiler/Core/Data/Stripped/InfoTable.hs +++ b/src/Juvix/Compiler/Core/Data/Stripped/InfoTable.hs @@ -46,6 +46,8 @@ data ConstructorInfo = ConstructorInfo _constructorTag :: Tag, _constructorType :: Type, _constructorArgNames :: [Maybe Text], + -- | _constructorArgsNum == length _constructorArgNames == length (typeArgs _constructorType) + _constructorArgsNum :: Int, _constructorFixity :: Maybe Fixity } diff --git a/src/Juvix/Compiler/Core/Translation/Stripped/FromCore.hs b/src/Juvix/Compiler/Core/Translation/Stripped/FromCore.hs index 39a1685e6e..2a0f238fba 100644 --- a/src/Juvix/Compiler/Core/Translation/Stripped/FromCore.hs +++ b/src/Juvix/Compiler/Core/Translation/Stripped/FromCore.hs @@ -144,6 +144,7 @@ translateConstructorInfo ConstructorInfo {..} = _constructorTag = _constructorTag, _constructorType = translateType _constructorType, _constructorArgNames, + _constructorArgsNum, _constructorFixity } diff --git a/src/Juvix/Compiler/Nockma/Translation/FromAsm.hs b/src/Juvix/Compiler/Nockma/Translation/FromAsm.hs index 30a6911d06..e158ad006a 100644 --- a/src/Juvix/Compiler/Nockma/Translation/FromAsm.hs +++ b/src/Juvix/Compiler/Nockma/Translation/FromAsm.hs @@ -6,8 +6,21 @@ import Juvix.Compiler.Nockma.Evaluator import Juvix.Compiler.Nockma.Pretty import Juvix.Compiler.Nockma.Stdlib import Juvix.Compiler.Pipeline.EntryPoint +import Juvix.Compiler.Tree.Language.Rep import Juvix.Prelude hiding (Atom, Path) +nockmaMemRep :: MemRep -> NockmaMemRep +nockmaMemRep = \case + MemRepTuple -> NockmaMemRepTuple + MemRepConstr -> NockmaMemRepConstr + MemRepTag -> NockmaMemRepConstr + MemRepUnit -> NockmaMemRepConstr + MemRepUnpacked {} -> NockmaMemRepConstr + +data NockmaMemRep + = NockmaMemRepConstr + | NockmaMemRepTuple + type UserFunctionId = Symbol data FunctionId @@ -41,11 +54,16 @@ data FunctionInfo = FunctionInfo data CompilerCtx = CompilerCtx { _compilerFunctionInfos :: HashMap FunctionId FunctionInfo, - _compilerConstructorArities :: ConstructorArities, + _compilerConstructorInfos :: ConstructorInfos, _compilerOptions :: CompilerOptions } -type ConstructorArities = HashMap Asm.Tag Natural +data ConstructorInfo = ConstructorInfo + { _constructorInfoArity :: Natural, + _constructorInfoMemRep :: NockmaMemRep + } + +type ConstructorInfos = HashMap Asm.Tag ConstructorInfo type Offset = Natural @@ -142,13 +160,20 @@ data Compiler m a where Save :: Bool -> m () -> Compiler m () CallStdlibOn :: StackId -> StdlibFunction -> Compiler m () AsmReturn :: Compiler m () - GetConstructorArity :: Asm.Tag -> Compiler m Natural + GetConstructorInfo :: Asm.Tag -> Compiler m ConstructorInfo GetFunctionArity :: FunctionId -> Compiler m Natural GetFunctionPath :: FunctionId -> Compiler m Path stackPath :: StackId -> Path stackPath s = indexStack (fromIntegral (fromEnum s)) +indexTuple :: Natural -> Natural -> Path +indexTuple len idx = + let lastL + | idx == len - 1 = [] + | otherwise = [L] + in replicate idx R ++ lastL + indexStack :: Natural -> Path indexStack idx = replicate idx R ++ [L] @@ -175,6 +200,7 @@ makeSem ''Compiler makeLenses ''CompilerOptions makeLenses ''CompilerFunction makeLenses ''CompilerCtx +makeLenses ''ConstructorInfo makeLenses ''FunctionInfo termFromParts :: (Bounded p, Enum p) => (p -> Term Natural) -> Term Natural @@ -206,8 +232,17 @@ fromAsmTable t = case t ^. Asm.infoMainFunction of fromAsm :: CompilerOptions -> Asm.Symbol -> Asm.InfoTable -> Cell Natural fromAsm opts mainSym Asm.InfoTable {..} = let funs = map compileFunction allFunctions - constrs :: ConstructorArities - constrs = fromIntegral . (^. Asm.constructorArgsNum) <$> _infoConstrs + mkConstructorInfo :: Asm.ConstructorInfo -> ConstructorInfo + mkConstructorInfo ci@Asm.ConstructorInfo {..} = + ConstructorInfo + { _constructorInfoArity = fromIntegral _constructorArgsNum, + _constructorInfoMemRep = nockmaMemRep (memRep ci (getInductiveInfo (ci ^. Asm.constructorInductive))) + } + constrs :: ConstructorInfos + constrs = mkConstructorInfo <$> _infoConstrs + + getInductiveInfo :: Symbol -> Asm.InductiveInfo + getInductiveInfo s = _infoInductives ^?! at s . _Just in runCompilerWith opts constrs funs mainFun where mainFun :: CompilerFunction @@ -235,20 +270,30 @@ fromAsmTable t = case t ^. Asm.infoMainFunction of _compilerFunction = compile _functionCode } + memRep :: Asm.ConstructorInfo -> Asm.InductiveInfo -> Asm.MemRep + memRep ci ind + | numArgs >= 1 && numConstrs == 1 = MemRepTuple + | otherwise = MemRepConstr + where + numConstrs = length (ind ^. Asm.inductiveConstructors) + numArgs = ci ^. Asm.constructorArgsNum + fromOffsetRef :: Asm.OffsetRef -> Natural fromOffsetRef = fromIntegral . (^. Asm.offsetRefOffset) -- | Generic constructors are encoded as [tag args], where args is a -- nil terminated list. -goConstructor :: Asm.Tag -> [Term Natural] -> Term Natural -goConstructor t args = case t of +goConstructor :: NockmaMemRep -> Asm.Tag -> [Term Natural] -> Term Natural +goConstructor mr t args = case t of Asm.BuiltinTag b -> makeConstructor $ \case ConstructorTag -> builtinTagToTerm b ConstructorArgs -> remakeList [] - Asm.UserTag tag -> - makeConstructor $ \case - ConstructorTag -> OpQuote # (fromIntegral (tag ^. Asm.tagUserWord) :: Natural) - ConstructorArgs -> remakeList args + Asm.UserTag tag -> case mr of + NockmaMemRepConstr -> + makeConstructor $ \case + ConstructorTag -> OpQuote # (fromIntegral (tag ^. Asm.tagUserWord) :: Natural) + ConstructorArgs -> remakeList args + NockmaMemRepTuple -> foldTerms (nonEmpty' args) compile :: forall r. (Members '[Compiler] r) => Asm.Code -> Sem r () compile = mapM_ goCommand @@ -303,6 +348,7 @@ compile = mapM_ goCommand Asm.DRef r -> pushDirectRef r Asm.ConstrRef r -> pushConstructorField + (r ^. Asm.fieldTag) (r ^. Asm.fieldRef) (fromIntegral (r ^. Asm.fieldOffset)) @@ -395,12 +441,28 @@ constVoid = makeConstructor $ \case ConstructorTag -> OpQuote # toNock (0 :: Natural) ConstructorArgs -> remakeList [] -pushConstructorFieldOnto :: (Members '[Compiler] r) => StackId -> Asm.DirectRef -> Natural -> Sem r () -pushConstructorFieldOnto s refToConstr argIx = - let path = directRefPath refToConstr ++ constructorPath ConstructorArgs ++ indexStack argIx - in pushOnto s (OpAddress # path) - -pushConstructorField :: (Members '[Compiler] r) => Asm.DirectRef -> Natural -> Sem r () +pushConstructorFieldOnto :: + (Members '[Compiler] r) => + StackId -> + Asm.Tag -> + Asm.DirectRef -> + Natural -> + Sem r () +pushConstructorFieldOnto s tag refToConstr argIx = do + info <- getConstructorInfo tag + let memrep = info ^. constructorInfoMemRep + arity = info ^. constructorInfoArity + path = case memrep of + NockmaMemRepConstr -> + directRefPath refToConstr + ++ constructorPath ConstructorArgs + ++ indexStack argIx + NockmaMemRepTuple -> + directRefPath refToConstr + ++ indexTuple arity argIx + pushOnto s (OpAddress # path) + +pushConstructorField :: (Members '[Compiler] r) => Asm.Tag -> Asm.DirectRef -> Natural -> Sem r () pushConstructorField = pushConstructorFieldOnto ValueStack directRefPath :: Asm.DirectRef -> Path @@ -438,9 +500,11 @@ closureArgsNum = do allocConstr :: (Members '[Compiler] r) => Asm.Tag -> Sem r () allocConstr tag = do - numArgs <- getConstructorArity tag - let args = [OpAddress # indexInStack ValueStack (pred i) | i <- [1 .. numArgs]] - constr = goConstructor tag args + info <- getConstructorInfo tag + let numArgs = info ^. constructorInfoArity + memrep = info ^. constructorInfoMemRep + args = [OpAddress # indexInStack ValueStack (pred i) | i <- [1 .. numArgs]] + constr = goConstructor memrep tag args pushOnto AuxStack constr popN numArgs moveTopFromTo AuxStack ValueStack @@ -515,7 +579,7 @@ runCompiler sem = do (ts, a) <- runOutputList (re sem) return (seqTerms ts, a) -runCompilerWith :: CompilerOptions -> ConstructorArities -> [CompilerFunction] -> CompilerFunction -> Cell Natural +runCompilerWith :: CompilerOptions -> ConstructorInfos -> [CompilerFunction] -> CompilerFunction -> Cell Natural runCompilerWith opts constrs libFuns mainFun = let entryCommand :: (Members '[Compiler] r) => Sem r () entryCommand = callFun (mainFun ^. compilerFunctionName) 0 @@ -547,7 +611,7 @@ runCompilerWith opts constrs libFuns mainFun = compilerCtx = CompilerCtx { _compilerFunctionInfos = functionInfos, - _compilerConstructorArities = constrs, + _compilerConstructorInfos = constrs, _compilerOptions = opts } @@ -811,6 +875,7 @@ constructorTagToTerm = \case Asm.BuiltinTag b -> builtinTagToTerm b caseCmd :: + forall r. (Members '[Compiler] r) => Maybe (Sem r ()) -> [(Asm.Tag, Sem r ())] -> @@ -818,11 +883,20 @@ caseCmd :: caseCmd defaultBranch = \case [] -> sequence_ defaultBranch (tag, b) : bs -> do - -- push the constructor tag at the top - push (OpAddress # topOfStack ValueStack ++ constructorPath ConstructorTag) - push (constructorTagToTerm tag) - testEq - branch b (caseCmd defaultBranch bs) + rep <- getConstructorMemRep tag + case rep of + NockmaMemRepConstr -> goRepConstr tag b bs + NockmaMemRepTuple + | null bs, isNothing defaultBranch -> b + | otherwise -> error "redundant branch. Impossible?" + where + goRepConstr :: Asm.Tag -> Sem r () -> [(Asm.Tag, Sem r ())] -> Sem r () + goRepConstr tag b bs = do + -- push the constructor tag at the top + push (OpAddress # topOfStack ValueStack ++ constructorPath ConstructorTag) + push (constructorTagToTerm tag) + testEq + branch b (caseCmd defaultBranch bs) branch' :: (Functor f, Members '[Output (Term Natural), Reader CompilerCtx] r) => @@ -837,8 +911,14 @@ branch' t f = do getFunctionArity' :: (Members '[Reader CompilerCtx] r) => FunctionId -> Sem r Natural getFunctionArity' s = asks (^?! compilerFunctionInfos . at s . _Just . functionInfoArity) -getConstructorArity' :: (Members '[Reader CompilerCtx] r) => Asm.Tag -> Sem r Natural -getConstructorArity' tag = asks (^?! compilerConstructorArities . at tag . _Just) +getConstructorInfo' :: (Members '[Reader CompilerCtx] r) => Asm.Tag -> Sem r ConstructorInfo +getConstructorInfo' tag = asks (^?! compilerConstructorInfos . at tag . _Just) + +getConstructorMemRep :: (Members '[Compiler] r) => Asm.Tag -> Sem r NockmaMemRep +getConstructorMemRep tag = (^. constructorInfoMemRep) <$> getConstructorInfo tag + +getConstructorArity :: (Members '[Compiler] r) => Asm.Tag -> Sem r Natural +getConstructorArity tag = (^. constructorInfoArity) <$> getConstructorInfo tag re :: (Member (Reader CompilerCtx) r) => Sem (Compiler ': r) a -> Sem (Output (Term Natural) ': r) a re = reinterpretH $ \case @@ -854,7 +934,7 @@ re = reinterpretH $ \case CallStdlibOn s f -> callStdlibOn' s f >>= pureT AsmReturn -> asmReturn' >>= pureT TestEqOn s -> testEqOn' s >>= pureT - GetConstructorArity s -> getConstructorArity' s >>= pureT + GetConstructorInfo s -> getConstructorInfo' s >>= pureT GetFunctionArity s -> getFunctionArity' s >>= pureT GetFunctionPath s -> getFunctionPath' s >>= pureT Crash -> outputT (OpAddress # OpAddress # OpAddress) @@ -1005,7 +1085,13 @@ pushNat = pushNatOnto ValueStack pushNatOnto :: (Member Compiler r) => StackId -> Natural -> Sem r () pushNatOnto s n = pushOnto s (OpQuote # toNock n) -compileAndRunNock' :: (Members '[Reader EvalOptions, Output (Term Natural)] r) => CompilerOptions -> ConstructorArities -> [CompilerFunction] -> CompilerFunction -> Sem r (Term Natural) +compileAndRunNock' :: + (Members '[Reader EvalOptions, Output (Term Natural)] r) => + CompilerOptions -> + ConstructorInfos -> + [CompilerFunction] -> + CompilerFunction -> + Sem r (Term Natural) compileAndRunNock' opts constrs funs mainfun = let Cell nockSubject t = runCompilerWith opts constrs funs mainfun in evalCompiledNock' nockSubject t diff --git a/src/Juvix/Compiler/Tree/Translation/FromCore.hs b/src/Juvix/Compiler/Tree/Translation/FromCore.hs index 0ef5f5507d..2839e21d83 100644 --- a/src/Juvix/Compiler/Tree/Translation/FromCore.hs +++ b/src/Juvix/Compiler/Tree/Translation/FromCore.hs @@ -329,8 +329,8 @@ translateConstructorInfo ci = { _constructorName = ci ^. Core.constructorName, _constructorLocation = ci ^. Core.constructorLocation, _constructorTag = ci ^. Core.constructorTag, - _constructorArgsNum = length (typeArgs ty), _constructorArgNames = ci ^. Core.constructorArgNames, + _constructorArgsNum = ci ^. Core.constructorArgsNum, _constructorType = ty, _constructorInductive = ci ^. Core.constructorInductive, _constructorRepresentation = MemRepConstr, diff --git a/test/Nockma/Compile/Positive.hs b/test/Nockma/Compile/Positive.hs index 57cf54ed03..1525a30195 100644 --- a/test/Nockma/Compile/Positive.hs +++ b/test/Nockma/Compile/Positive.hs @@ -91,28 +91,41 @@ functionCode = \case data ConstructorName = ConstructorFalse | ConstructorTrue - | ConstructorWrapper + | ConstructorTagged | ConstructorPair + | ConstructorTuple deriving stock (Eq, Bounded, Enum) constructorTag :: ConstructorName -> Asm.Tag constructorTag n = Asm.UserTag (Asm.TagUser defaultModuleId (fromIntegral (fromEnum n))) -constructorArity :: ConstructorName -> Natural -constructorArity = \case - ConstructorFalse -> 0 - ConstructorTrue -> 0 - ConstructorWrapper -> 1 - ConstructorPair -> 2 +constructorInfo :: ConstructorName -> ConstructorInfo +constructorInfo = \case + ConstructorFalse -> defaultInfo 0 + ConstructorTrue -> defaultInfo 0 + ConstructorTagged -> defaultInfo 1 + ConstructorPair -> defaultInfo 2 + ConstructorTuple -> + ConstructorInfo + { _constructorInfoArity = 2, + _constructorInfoMemRep = NockmaMemRepTuple + } -exampleConstructors :: ConstructorArities +defaultInfo :: Natural -> ConstructorInfo +defaultInfo ari = + ConstructorInfo + { _constructorInfoArity = ari, + _constructorInfoMemRep = NockmaMemRepConstr + } + +exampleConstructors :: ConstructorInfos exampleConstructors = hashMap $ - [ (constructorTag n, constructorArity n) + [ (constructorTag n, constructorInfo n) | n <- allElements ] - ++ [ (Asm.BuiltinTag Asm.TagTrue, 0), - (Asm.BuiltinTag Asm.TagFalse, 0) + ++ [ (Asm.BuiltinTag Asm.TagTrue, defaultInfo 0), + (Asm.BuiltinTag Asm.TagFalse, defaultInfo 0) ] exampleFunctions :: [CompilerFunction] @@ -367,7 +380,7 @@ tests = allocConstr (constructorTag ConstructorFalse), defTest "alloc unary constructor" (eqStack ValueStack [nock| [[2 [[55 66] nil] nil] nil]|]) $ do push (OpQuote # (55 :: Natural) # (66 :: Natural)) - allocConstr (constructorTag ConstructorWrapper), + allocConstr (constructorTag ConstructorTagged), defTest "alloc binary constructor" (eqStack ValueStack [nock| [[3 [9 7 nil] nil] nil] |]) $ do pushNat 7 pushNat 9 @@ -465,31 +478,54 @@ tests = (eqStack ValueStack [nock| [777 [2 [123 nil] nil] nil] |]) $ do pushNat 123 - allocConstr (constructorTag ConstructorWrapper) + allocConstr (constructorTag ConstructorTagged) caseCmd Nothing - [ (constructorTag ConstructorWrapper, pushNat 777) + [ (constructorTag ConstructorTagged, pushNat 777) ], defTest "cmdCase: default branch" (eqStack ValueStack [nock| [5 nil] |]) $ do pushNat 123 - allocConstr (constructorTag ConstructorWrapper) + allocConstr (constructorTag ConstructorTagged) caseCmd (Just (pop >> pushNat 5)) [ (constructorTag ConstructorFalse, pushNat 777) ], + defTest + "cmdCase: case on pair (NockmaMemRepTuple)" + (eqStack ValueStack [nock| [[70 50 10] nil] |]) + $ do + let t = constructorTag ConstructorTuple + pushNat 10 + pushNat 50 + allocConstr t + caseCmd + Nothing + [ ( t, + do + pushConstructorFieldOnto AuxStack t Asm.StackRef 0 + pushConstructorFieldOnto AuxStack t Asm.StackRef 1 + pushConstructorFieldOnto AuxStack t Asm.StackRef 1 + moveTopFromTo AuxStack ValueStack + moveTopFromTo AuxStack ValueStack + moveTopFromTo AuxStack ValueStack + add + add + allocConstr t + ) + ], defTest "cmdCase: second branch" (eqStack ValueStack [nock| [5 nil] |]) $ do pushNat 123 - allocConstr (constructorTag ConstructorWrapper) + allocConstr (constructorTag ConstructorTagged) caseCmd (Just (pushNat 0)) [ (constructorTag ConstructorFalse, pushNat 0), - (constructorTag ConstructorWrapper, pop >> pushNat 5) + (constructorTag ConstructorTagged, pop >> pushNat 5) ], defTest "cmdCase: case on builtin true" @@ -518,8 +554,8 @@ tests = pushNat 10 pushNat 20 allocConstr (constructorTag ConstructorPair) - pushConstructorFieldOnto TempStack Asm.StackRef 0 - pushConstructorFieldOnto TempStack Asm.StackRef 1 + pushConstructorFieldOnto TempStack (constructorTag ConstructorPair) Asm.StackRef 0 + pushConstructorFieldOnto TempStack (constructorTag ConstructorPair) Asm.StackRef 1 addOn TempStack, defTest "trace"