Compile: More precise monads in genStructs

1 files changed, 13 insertions, 11 deletions

diff --git a/src/Compile.hs b/src/Compile.hs
index 48a03b5..582a1df 100644
--- a/src/Compile.hs
+++ b/src/Compile.hs
@@ -4,8 +4,10 @@
 {-# LANGUAGE TypeApplications #-}
 module Compile (compile) where
 
+import Control.Monad.Trans.Class (lift)
 import Control.Monad.Trans.State.Strict
-import Data.Bifunctor (first, second)
+import Control.Monad.Trans.Writer.CPS
+import Data.Bifunctor (first)
 import Data.Foldable (toList)
 import Data.Functor.Const
 import qualified Data.Functor.Product as Product
@@ -190,18 +192,18 @@ genStruct name topty = case topty of
   where
     com = ppTy 0 topty
 
--- State: (already-generated (skippable) struct names, the structs in declaration order)
-genStructs :: Ty -> State (Set String, Bag StructDecl) ()
+-- State: already-generated (skippable) struct names
+-- Writer: the structs in declaration order
+genStructs :: Ty -> WriterT (Bag StructDecl) (State (Set String)) ()
 genStructs ty = do
   let name = genStructName ty
-  seen <- gets ((name `Set.member`) . fst)
+  seen <- lift $ gets (name `Set.member`)
 
-  case (if seen then Nothing else genStruct name ty) of
-    Nothing -> pure ()
-
-    Just decl -> do
+  if seen
+    then pure ()
+    else do
       -- already mark this struct as generated now, so we don't generate it twice
-      modify (first (Set.insert name))
+      lift $ modify (Set.insert name)
 
       case ty of
         TNil -> pure ()
@@ -212,10 +214,10 @@ genStructs ty = do
         TScal _ -> pure ()
         TAccum t -> genStructs t
 
-      modify (second (<> pure decl))
+      tell (maybe mempty pure (genStruct name ty))
 
 genAllStructs :: Foldable t => t Ty -> [StructDecl]
-genAllStructs tys = toList . snd $ execState (mapM_ genStructs tys) (mempty, mempty)
+genAllStructs tys = toList $ evalState (execWriterT (mapM_ genStructs tys)) mempty
 
 data CompState = CompState
   { csStructs :: Set Ty