Skip to content

Commit

Permalink
Support MemRepTuple in the Nockma backend (#2586)
Browse files Browse the repository at this point in the history 
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 <git@paulcadman.dev>
  • Loading branch information
@janmasrovira @paulcadman
janmasrovira and paulcadman authored Jan 23, 2024
1 parent 8005089 commit 510490a
Show file tree
Hide file tree
Showing 5 changed files with 176 additions and 51 deletions.
2 changes: 2 additions & 0 deletions src/Juvix/Compiler/Core/Data/Stripped/InfoTable.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down
1 change: 1 addition & 0 deletions src/Juvix/Compiler/Core/Translation/Stripped/FromCore.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,7 @@ translateConstructorInfo ConstructorInfo {..} =
_constructorTag = _constructorTag,
_constructorType = translateType _constructorType,
_constructorArgNames,
_constructorArgsNum,
_constructorFixity
}

Expand Down
148 changes: 117 additions & 31 deletions src/Juvix/Compiler/Nockma/Translation/FromAsm.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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]

Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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))

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -547,7 +611,7 @@ runCompilerWith opts constrs libFuns mainFun =
compilerCtx =
CompilerCtx
{ _compilerFunctionInfos = functionInfos,
_compilerConstructorArities = constrs,
_compilerConstructorInfos = constrs,
_compilerOptions = opts
}

Expand Down Expand Up @@ -811,18 +875,28 @@ constructorTagToTerm = \case
Asm.BuiltinTag b -> builtinTagToTerm b

caseCmd ::
forall r.
(Members '[Compiler] r) =>
Maybe (Sem r ()) ->
[(Asm.Tag, Sem r ())] ->
Sem r ()
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) =>
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/Juvix/Compiler/Tree/Translation/FromCore.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading

0 comments on commit 510490a

Please sign in to comment.