Compile: Improve/simplify struct generation

1 files changed, 68 insertions, 57 deletions

diff --git a/src/Compile.hs b/src/Compile.hs
index 064e0b6..0ab7ea4 100644
--- a/src/Compile.hs
+++ b/src/Compile.hs
@@ -126,7 +126,7 @@ data CExpr
   | CECall String [CExpr]  -- ^ function(arg1, ..., argn)
   | CEBinop CExpr String CExpr  -- ^ expr + expr
   | CEIf CExpr CExpr CExpr  -- ^ expr ? expr : expr
-  | CECast String CExpr  -- ^ (<type)<expr>
+  | CECast String CExpr  -- ^ (<type>)<expr>
   deriving (Show)
 
 printStructDecl :: StructDecl -> ShowS
@@ -215,23 +215,31 @@ repSTy (STScal st) = case st of
   STBool -> "uint8_t"
 repSTy t = genStructName t
 
-genStructName :: STy t -> String
-genStructName = \t -> "ty_" ++ gen t where
-  -- all tags start with a letter, so the array mangling is unambiguous.
-  gen :: STy t -> String
-  gen STNil = "n"
-  gen (STPair a b) = 'P' : gen a ++ gen b
-  gen (STEither a b) = 'E' : gen a ++ gen b
-  gen (STLEither a b) = 'L' : gen a ++ gen b
-  gen (STMaybe t) = 'M' : gen t
-  gen (STArr n t) = "A" ++ show (fromSNat n) ++ gen t
-  gen (STScal st) = case st of
-    STI32 -> "i"
-    STI64 -> "j"
-    STF32 -> "f"
-    STF64 -> "d"
-    STBool -> "b"
-  gen (STAccum t) = 'C' : gen (fromSMTy t)
+genStructName, genArrBufStructName :: STy t -> String
+(genStructName, genArrBufStructName) =
+  (\t -> "ty_" ++ gen t
+  ,\case STArr _ t -> "ty_A_" ++ gen t ++ "_buf"  -- just like the normal type, but with _ for the dimension
+         t -> error $ "genArrBufStructName: not an array type: " ++ show t)
+  where
+    -- all tags start with a letter, so the array mangling is unambiguous.
+    gen :: STy t -> String
+    gen STNil = "n"
+    gen (STPair a b) = 'P' : gen a ++ gen b
+    gen (STEither a b) = 'E' : gen a ++ gen b
+    gen (STLEither a b) = 'L' : gen a ++ gen b
+    gen (STMaybe t) = 'M' : gen t
+    gen (STArr n t) = "A" ++ show (fromSNat n) ++ gen t
+    gen (STScal st) = case st of
+      STI32 -> "i"
+      STI64 -> "j"
+      STF32 -> "f"
+      STF64 -> "d"
+      STBool -> "b"
+    gen (STAccum t) = 'C' : gen (fromSMTy t)
+
+-- The subtrees contain structs used in the bodies of the structs in this node.
+data StructTree = TreeNode [StructDecl] [StructTree]
+  deriving (Show)
 
 -- | This function generates the actual struct declarations for each of the
 -- types in our language. It thus implicitly "documents" the layout of the
@@ -239,59 +247,56 @@ genStructName = \t -> "ty_" ++ gen t where
 --
 -- For accumulation it is important that for struct representations of monoid
 -- types, the all-zero-bytes value corresponds to the zero value of that type.
-genStruct :: String -> STy t -> [StructDecl]
-genStruct name topty = case topty of
+buildStructTree :: STy t -> StructTree
+buildStructTree topty = case topty of
   STNil ->
-    [StructDecl name "" com]
+    TreeNode [StructDecl name "" com] []
   STPair a b ->
-    [StructDecl name (repSTy a ++ " a; " ++ repSTy b ++ " b;") com]
+    TreeNode [StructDecl name (repSTy a ++ " a; " ++ repSTy b ++ " b;") com]
+             [buildStructTree a, buildStructTree b]
   STEither a b ->  -- 0 -> l, 1 -> r
-    [StructDecl name ("uint8_t tag; union { " ++ repSTy a ++ " l; " ++ repSTy b ++ " r; };") com]
+    TreeNode [StructDecl name ("uint8_t tag; union { " ++ repSTy a ++ " l; " ++ repSTy b ++ " r; };") com]
+             [buildStructTree a, buildStructTree b]
   STLEither a b ->  -- 0 -> nil, 1 -> l, 2 -> r
-    [StructDecl name ("uint8_t tag; union { " ++ repSTy a ++ " l; " ++ repSTy b ++ " r; };") com]
+    TreeNode [StructDecl name ("uint8_t tag; union { " ++ repSTy a ++ " l; " ++ repSTy b ++ " r; };") com]
+             [buildStructTree a, buildStructTree b]
   STMaybe t ->  -- 0 -> nothing, 1 -> just
-    [StructDecl name ("uint8_t tag; " ++ repSTy t ++ " j;") com]
+    TreeNode [StructDecl name ("uint8_t tag; " ++ repSTy t ++ " j;") com]
+             [buildStructTree t]
   STArr n t ->
     -- The buffer is trailed by a VLA for the actual array data.
     -- TODO: no buffer if n = 0
-    [StructDecl (name ++ "_buf") ("size_t refc; " ++ repSTy t ++ " xs[];") ""
-    ,StructDecl name (name ++ "_buf *buf; size_t sh[" ++ show (fromSNat n) ++ "];") com]
+    TreeNode [StructDecl (genArrBufStructName topty) ("size_t refc; " ++ repSTy t ++ " xs[];") ""
+             ,StructDecl name (genArrBufStructName topty ++ " *buf; size_t sh[" ++ show (fromSNat n) ++ "];") com]
+             [buildStructTree t]
   STScal _ ->
-    []
+    TreeNode [] []
   STAccum t ->
-    [StructDecl (name ++ "_buf") (repSTy (fromSMTy t) ++ " ac;") ""
-    ,StructDecl name (name ++ "_buf *buf;") com]
+    TreeNode [StructDecl (name ++ "_buf") (repSTy (fromSMTy t) ++ " ac;") ""
+             ,StructDecl name (name ++ "_buf *buf;") com]
+             [buildStructTree (fromSMTy t)]
   where
+    name = genStructName topty
     com = ppSTy 0 topty
 
 -- State: already-generated (skippable) struct names
 -- Writer: the structs in declaration order
-genStructs :: STy t -> WriterT (Bag StructDecl) (State (Set String)) ()
-genStructs ty = do
-  let name = genStructName ty
-  seen <- lift $ gets (name `Set.member`)
-
-  if seen
-    then pure ()
-    else do
-      -- already mark this struct as generated now, so we don't generate it
-      -- twice (unnecessary because no recursive types, but y'know)
-      lift $ modify (Set.insert name)
-
-      () <- case ty of
-        STNil -> pure ()
-        STPair a b -> genStructs a >> genStructs b
-        STEither a b -> genStructs a >> genStructs b
-        STLEither a b -> genStructs a >> genStructs b
-        STMaybe t -> genStructs t
-        STArr _ t -> genStructs t
-        STScal _ -> pure ()
-        STAccum t -> genStructs (fromSMTy t)
-
-      tell (BList (genStruct name ty))
+genStructTreeW :: StructTree -> WriterT (Bag StructDecl) (State (Set String)) ()
+genStructTreeW (TreeNode these deps) = do
+  seen <- lift get
+  case filter ((`Set.notMember` seen) . nameOf) these of
+    [] -> pure ()
+    structs -> do
+      lift $ modify (Set.fromList (map nameOf structs) <>)
+      mapM_ genStructTreeW deps
+      tell (BList structs)
+  where
+    nameOf (StructDecl name _ _) = name
 
 genAllStructs :: Foldable t => t (Some STy) -> [StructDecl]
-genAllStructs tys = toList $ evalState (execWriterT (mapM_ (\(Some t) -> genStructs t) tys)) mempty
+genAllStructs tys =
+  let m = mapM_ (\(Some t) -> genStructTreeW (buildStructTree t)) tys
+  in toList (evalState (execWriterT m) mempty)
 
 data CompState = CompState
   { csStructs :: Set (Some STy)
@@ -340,6 +345,12 @@ emitStruct ty = CompM $ do
   modify $ \s -> s { csStructs = Set.insert (Some ty) (csStructs s) }
   return (genStructName ty)
 
+-- | Also returns the name of the array buffer struct
+emitArrStruct :: STy t -> CompM (String, String)
+emitArrStruct ty = CompM $ do
+  modify $ \s -> s { csStructs = Set.insert (Some ty) (csStructs s) }
+  return (genStructName ty, genArrBufStructName ty)
+
 emitTLD :: String -> CompM ()
 emitTLD decl = CompM $ modify $ \s -> s { csTopLevelDecls = csTopLevelDecls s <> pure decl }
 
@@ -745,13 +756,13 @@ compile' env = \case
     return (CELit retvar)
 
   EConstArr _ n t (Array sh vec) -> do
-    strname <- emitStruct (STArr n (STScal t))
+    (strname, bufstrname) <- emitArrStruct (STArr n (STScal t))
     tldname <- genName' "carraybuf"
     -- Give it a refcount of _half_ the size_t max, so that it can be
     -- incremented and decremented at will and will "never" reach anything
     -- where something happens
-    emitTLD $ "static " ++ strname ++ "_buf " ++ tldname ++ " = " ++
-              "(" ++ strname ++ "_buf){.refc = (size_t)1<<63, " ++
+    emitTLD $ "static " ++ bufstrname ++ " " ++ tldname ++ " = " ++
+              "(" ++ bufstrname ++ "){.refc = (size_t)1<<63, " ++
               ".xs = {" ++ intercalate "," (map (compileScal False t) (toList vec)) ++ "}};"
     return (CEStruct strname
               [("buf", CEAddrOf (CELit tldname))