Split the Ast file into multiple sub files.

Make some more changes to the generic implementation of EasySwitchStage.

2 files changed, 121 insertions, 115 deletions

diff --git a/src/Language/Fiddle/Compiler/Stage1.hs b/src/Language/Fiddle/Compiler/Stage1.hs
index a17afa1..d2fe885 100644
--- a/src/Language/Fiddle/Compiler/Stage1.hs
+++ b/src/Language/Fiddle/Compiler/Stage1.hs
@@ -5,11 +5,11 @@
 
 module Language.Fiddle.Compiler.Stage1 (toStage2) where
 
-import qualified Data.List.NonEmpty as NonEmpty
 import Control.Monad.Identity (Identity (..))
 import Control.Monad.State (get, gets, modify, put)
 import qualified Data.Char as Char
 import Data.List (intercalate)
+import qualified Data.List.NonEmpty as NonEmpty
 import Data.Text (Text)
 import qualified Data.Text as Text
 import Data.Type.Bool
@@ -69,9 +69,9 @@ toStage2 (FiddleUnit decls annot) = do
       FiddleUnit <$> reconfigureFiddleDecls (Path []) decls <*> pure annot
   return a
 
-reconfigureFiddleDecls :: Path -> [Directed (FiddleDecl Stage1) I Annot] -> M Annot [Directed (FiddleDecl Stage2) I Annot]
+reconfigureFiddleDecls :: Path -> [Directed FiddleDecl Stage1 I Annot] -> M Annot [Directed FiddleDecl Stage2 I Annot]
 reconfigureFiddleDecls p decls = do
-  -- (Stage2CompilerState anonymousObjTypes anonymousBitsTypes, decls) <- pushState $ do
+  -- Stage2CompilerState anonymousObjTypes anonymousBitsTypes, decls <- pushState $ do
   --   put (Stage2CompilerState [] [])
   --   gets (,) <*> mapM (fiddleDeclToStage2 p) decls
 
@@ -95,6 +95,7 @@ reconfigureFiddleDecls p decls = do
 pushId :: Identifier f a -> Path -> Path
 pushId (Identifier str _) (Path lst) =
   Path (PathExpression (Text.unpack str) : lst)
+
 pushName :: Name f a -> Path -> Path
 pushName (Name idents _) path =
   foldl (flip pushId) path idents
@@ -134,7 +135,10 @@ objTypeBodyToStage2 path (ObjTypeBody bodyType decls annot) =
 
 objTypeDeclToStage2 :: Path -> ObjTypeDecl Stage1 I Annot -> M Annot (ObjTypeDecl Stage2 I Annot)
 objTypeDeclToStage2 path = \case
-  (AssertPosStatement expr annot) -> AssertPosStatement <$> toStage2Expr expr <*> pure annot
+  (AssertPosStatement w expr annot) ->
+    AssertPosStatement w
+      <$> toStage2Expr expr
+      <*> pure annot
   (TypeSubStructure (Identity deferredBody) maybeIdent annot) ->
     let path' = maybe path (`pushId` path) maybeIdent
      in TypeSubStructure . Identity
@@ -185,7 +189,7 @@ registerBitsTypeRefToStage2 path = \case
       <*> pure annot
   RegisterBitsReference name annot -> return (RegisterBitsReference name annot)
   RegisterBitsJustBits expr annot -> RegisterBitsJustBits <$> toStage2Expr expr <*> pure annot
-  RegisterBitsAnonymousType anonType annot -> do
+  RegisterBitsAnonymousType _ anonType annot -> do
     ident <- internAnonymousBitsType path =<< anonymousBitsTypeToStage2 path anonType
     return $ RegisterBitsReference (identToName ident) annot
 
@@ -202,7 +206,7 @@ anonymousBitsTypeToStage2 path = \case
 
 objectTypeToStage2 :: Path -> ObjType Stage1 I Annot -> M Annot (ObjType Stage2 I Annot)
 objectTypeToStage2 path = \case
-  (AnonymousObjType (Identity body) annot) -> do
+  (AnonymousObjType _ (Identity body) annot) -> do
     body' <- objTypeBodyToStage2 path body
     identifier <- internObjType path body'
     return (ReferencedObjType (identToName identifier) annot)
diff --git a/src/Language/Fiddle/Compiler/Stage2.hs b/src/Language/Fiddle/Compiler/Stage2.hs
index 57b0b55..2035e3d 100644
--- a/src/Language/Fiddle/Compiler/Stage2.hs
+++ b/src/Language/Fiddle/Compiler/Stage2.hs
@@ -5,11 +5,11 @@
 
 -- Stage3 doesn't change much from Stage2. Stage3 primarily removes the assert
 -- statements and checks that they are consistent with the calculations.
-module Language.Fiddle.Compiler.Stage2 where
+module Language.Fiddle.Compiler.Stage2 (toStage3) where
 
 import Control.Monad (forM, forM_, unless, when)
 import Control.Monad.RWS (MonadState (get), MonadWriter (tell), gets, modify')
-import Data.Foldable (foldlM, Foldable (toList))
+import Data.Foldable (Foldable (toList), foldlM)
 import Data.Functor.Identity
 import qualified Data.IntMap as IntMap
 import Data.Kind (Type)
@@ -110,12 +110,6 @@ insertTypeSize (Identifier s _) size = do
                 insertScope fullName (Right size) (inScope stage3State)
             }
 
--- addTypeSize :: Identifier f a -> SizeBits -> Compile Stage3State ()
--- addTypeSize (Identifier s _) size = do
---   modify' $
---     \stage3State ->
---       stage3State {typeSizes = Map.insert (Text.unpack s) size (typeSizes stage3State)}
-
 lookupTypeSize :: Name I Annot -> Compile Stage3State SizeBits
 lookupTypeSize (Name idents a) = do
   let path = fmap (\(Identifier s _) -> Text.unpack s) idents
@@ -135,11 +129,6 @@ lookupTypeSize (Name idents a) = do
         ]
       compilationFailure
 
-expressionToStage3 :: Expression Stage2 f Annot -> Expression Stage3 f Annot
-expressionToStage3 = \case
-  LitNum n a -> LitNum n a
-  Var i a -> Var i a
-
 emptyState = Stage3State mempty mempty
 
 toStage3 :: FiddleUnit Stage2 I Annot -> Compile () (FiddleUnit Stage3 I Annot)
@@ -178,7 +167,10 @@ getTypeSize (EnumBitType expr (Identity (EnumBody constants _)) ann) = do
               tell
                 [ Diagnostic
                     Error
-                    (printf "Enum constant too large. Max allowed %d\n" ((2 :: Int) ^ declaredSize))
+                    ( printf
+                        "Enum constant too large. Max allowed %d\n"
+                        ((2 :: Int) ^ declaredSize)
+                    )
                     (unCommented (annot enumConst))
                 ]
 
@@ -186,7 +178,8 @@ getTypeSize (EnumBitType expr (Identity (EnumBody constants _)) ann) = do
         )
         IntMap.empty
         constants
-    let missing = filter (not . (`IntMap.member` imap)) [0 .. 2 ^ declaredSize - 1]
+    let missing =
+          filter (not . (`IntMap.member` imap)) [0 .. 2 ^ declaredSize - 1]
     unless (null missing) $
       tell
         [ Diagnostic
@@ -206,7 +199,8 @@ addCurrentScope s = do
   modify' $ \st@(Stage3State {scopePath = (ScopePath current others)}) ->
     st {scopePath = ScopePath (current ++ s) others}
 
-fiddleDeclToStage3 :: FiddleDecl Stage2 I Annot -> Compile Stage3State (FiddleDecl Stage3 I Annot)
+fiddleDeclToStage3 ::
+  FiddleDecl Stage2 I Annot -> Compile Stage3State (FiddleDecl Stage3 I Annot)
 fiddleDeclToStage3 = \case
   OptionDecl i1 i2 a -> return $ OptionDecl i1 i2 a
   PackageDecl n@(Name idents _) body a -> do
@@ -227,97 +221,120 @@ fiddleDeclToStage3 = \case
                 }
       )
     return $ UsingDecl n a
-  LocationDecl id expr a -> return $ LocationDecl id (expressionToStage3 expr) a
+  LocationDecl id expr a -> return $ LocationDecl id (switchStage expr) a
   BitsDecl id typ a -> do
     typeSize <- getTypeSize typ
     insertTypeSize id typeSize
-    BitsDecl id <$> bitTypeToStage3 typ <*> pure a
+    return $ BitsDecl id (switchStage typ) a
   ObjTypeDecl ident body a ->
-    ObjTypeDecl ident <$> mapM (\bt -> fst <$> objTypeBodyToStage3 bt 0) body <*> pure a
+    ObjTypeDecl ident
+      <$> mapM (\bt -> fst <$> objTypeBodyToStage3 bt 0) body
+      <*> pure a
   ImportDecl importStatement a -> return $ ImportDecl importStatement a
   ObjectDecl ident expr typ a ->
-    ObjectDecl ident (expressionToStage3 expr) <$> objTypeToStage3 typ <*> pure a
+    ObjectDecl
+      ident
+      (switchStage expr)
+      <$> objTypeToStage3 typ
+      <*> pure a
 
-objTypeToStage3 :: ObjType Stage2 I Annot -> Compile Stage3State (ObjType Stage3 I Annot)
+objTypeToStage3 ::
+  ObjType Stage2 I Annot -> Compile Stage3State (ObjType Stage3 I Annot)
 objTypeToStage3 = \case
   ArrayObjType objtype expr a ->
     ArrayObjType
       <$> objTypeToStage3 objtype
-      <*> pure (expressionToStage3 expr)
+      <*> pure (switchStage expr)
       <*> pure a
   ReferencedObjType ident a -> return $ ReferencedObjType ident a
 
 registerBodyToStage3 ::
   RegisterBody Stage2 I Annot ->
   Compile Stage3State (RegisterBody Stage3 I Annot, Word32)
-registerBodyToStage3 (RegisterBody bodyType (Identity deferredRegisterBody) a') = do
-  let isUnion = case bodyType of
-        Union {} -> True
-        _ -> False
-
-  case deferredRegisterBody of
-    DeferredRegisterBody decls a -> do
-      (cur, returned) <-
-        foldlM
-          ( \(cursor, returned) decl ->
-              case undirected decl of
-                ReservedBits expr a -> do
-                  size <- fromIntegral <$> exprToSize expr
-                  let s3 = mapDirected (const $ ReservedBits (expressionToStage3 expr) a) decl
-                  if isUnion
-                    then checkUnion cursor size (s3 : returned) a
-                    else
-                      return (cursor + size, s3 : returned)
-                BitsSubStructure registerBody maybeIdent annot -> do
-                  checkBitsSubStructure registerBody maybeIdent annot
-
-                  (newBody, subsize) <- registerBodyToStage3 registerBody
-                  let s3 = mapDirected (const $ BitsSubStructure newBody maybeIdent annot) decl
-
-                  if isUnion
-                    then checkUnion cursor subsize (s3 : returned) a
-                    else
-                      return (cursor + subsize, s3 : returned)
-                DefinedBits modifier identifier typeref a -> do
-                  (s3TypeRef, size) <- registerBitsTypeRefToStage3 typeref
-                  let s3 = mapDirected (const $ DefinedBits modifier identifier s3TypeRef a) decl
-
-                  if isUnion
-                    then checkUnion cursor size (s3 : returned) a
-                    else
-                      return (cursor + size, s3 : returned)
+registerBodyToStage3
+  (RegisterBody bodyType (Identity deferredRegisterBody) a') = do
+    let isUnion = case bodyType of
+          Union {} -> True
+          _ -> False
+
+    case deferredRegisterBody of
+      DeferredRegisterBody decls a -> do
+        (cur, returned) <-
+          foldlM
+            ( \(cursor, returned) decl ->
+                case undirected decl of
+                  ReservedBits expr a -> do
+                    size <- fromIntegral <$> exprToSize expr
+                    let s3 =
+                          mapDirected
+                            (const $ ReservedBits (switchStage expr) a)
+                            decl
+                    if isUnion
+                      then checkUnion cursor size (s3 : returned) a
+                      else
+                        return (cursor + size, s3 : returned)
+                  BitsSubStructure registerBody maybeIdent annot -> do
+                    checkBitsSubStructure registerBody maybeIdent annot
+
+                    (newBody, subsize) <- registerBodyToStage3 registerBody
+                    let s3 =
+                          mapDirected
+                            (const $ BitsSubStructure newBody maybeIdent annot)
+                            decl
+
+                    if isUnion
+                      then checkUnion cursor subsize (s3 : returned) a
+                      else
+                        return (cursor + subsize, s3 : returned)
+                  DefinedBits modifier identifier typeref a -> do
+                    (s3TypeRef, size) <- registerBitsTypeRefToStage3 typeref
+                    let s3 =
+                          mapDirected
+                            (const $ DefinedBits modifier identifier s3TypeRef a)
+                            decl
+
+                    if isUnion
+                      then checkUnion cursor size (s3 : returned) a
+                      else
+                        return (cursor + size, s3 : returned)
+            )
+            (0, [])
+            decls
+
+        return
+          ( RegisterBody
+              bodyType
+              (Identity (DeferredRegisterBody (reverse returned) a))
+              a',
+            cur
           )
-          (0, [])
-          decls
-
-      return (RegisterBody bodyType (Identity (DeferredRegisterBody (reverse returned) a)) a', cur)
-  where
-    checkBitsSubStructure
-      (RegisterBody bodyType (Identity (DeferredRegisterBody decls _)) _)
-      maybeIdent
-      annot =
-        let emitWarning s = tell [Diagnostic Warning s (unCommented annot)]
-         in case () of
-              ()
-                | [_] <- decls,
-                  (Union {}) <- bodyType ->
-                    emitWarning "Union with a single field. Should this be a struct?"
-              ()
-                | [_] <- decls,
-                  (Struct {}) <- bodyType,
-                  Nothing <- maybeIdent ->
-                    emitWarning "Anonymous sub-struct with single field is superfluous."
-              ()
-                | [] <- decls ->
-                    emitWarning
-                      ( printf
-                          "Empty sub-%s is superfluous."
-                          ( case bodyType of
-                              Union {} -> "union"
-                              Struct {} -> "struct"
-                          )
-                      )
-              _ -> return ()
+    where
+      checkBitsSubStructure
+        (RegisterBody bodyType (Identity (DeferredRegisterBody decls _)) _)
+        maybeIdent
+        annot =
+          let emitWarning s = tell [Diagnostic Warning s (unCommented annot)]
+           in case () of
+                ()
+                  | [_] <- decls,
+                    (Union {}) <- bodyType ->
+                      emitWarning "Union with a single field. Should this be a struct?"
+                ()
+                  | [_] <- decls,
+                    (Struct {}) <- bodyType,
+                    Nothing <- maybeIdent ->
+                      emitWarning "Anonymous sub-struct with single field is superfluous."
+                ()
+                  | [] <- decls ->
+                      emitWarning
+                        ( printf
+                            "Empty sub-%s is superfluous."
+                            ( case bodyType of
+                                Union {} -> "union"
+                                Struct {} -> "struct"
+                            )
+                        )
+                _ -> return ()
 
 checkUnion :: Word32 -> Word32 -> b -> Commented SourceSpan -> Compile Stage3State (Word32, b)
 checkUnion cursor subsize ret a = do
@@ -344,13 +361,13 @@ registerBitsTypeRefToStage3 = \case
     (ref', size) <- registerBitsTypeRefToStage3 ref
     multiplier <- exprToSize expr
     return
-      ( RegisterBitsArray ref' (expressionToStage3 expr) a,
+      ( RegisterBitsArray ref' (switchStage expr) a,
         size * fromIntegral multiplier
       )
   RegisterBitsReference name a ->
     (RegisterBitsReference name a,) <$> lookupTypeSize name
   RegisterBitsJustBits expr a ->
-    (RegisterBitsJustBits (expressionToStage3 expr) a,)
+    (RegisterBitsJustBits (switchStage expr) a,)
       . fromIntegral
       <$> exprToSize expr
 
@@ -374,7 +391,7 @@ objTypeBodyToStage3 (ObjTypeBody bodyType decls a) startOff = do
                           RegisterDecl
                             mMod
                             mIdent
-                            (expressionToStage3 expr)
+                            (switchStage expr)
                             s3RegisterBody
                             a
                       )
@@ -440,13 +457,13 @@ objTypeBodyToStage3 (ObjTypeBody bodyType decls a) startOff = do
                       (unCommented a)
                   ]
               let size = size' `div` 8
-              let s3 = mapDirected (const $ ReservedDecl (expressionToStage3 expr) annot) decl
+              let s3 = mapDirected (const $ ReservedDecl (switchStage expr) annot) decl
               if isUnion
                 then
                   checkUnion cursor size (s3 : returned) a
                 else
                   return (cursor + size, s3 : returned)
-            AssertPosStatement expr a -> do
+            AssertPosStatement _ expr a -> do
               declaredPos <- fromIntegral <$> exprToSize expr
 
               let expectedPos = if isUnion then startOff else cursor + startOff
@@ -500,21 +517,6 @@ objTypeBodyToStage3 (ObjTypeBody bodyType decls a) startOff = do
     pushApply (Just (a, b)) = (Just a, Just b)
     pushApply Nothing = (Nothing, Nothing)
 
-bitTypeToStage3 :: BitType Stage2 I Annot -> Compile Stage3State (BitType Stage3 I Annot)
-bitTypeToStage3 (EnumBitType expr body a) =
-  EnumBitType (expressionToStage3 expr)
-    <$> mapM enumBodyToStage3 body
-    <*> pure a
-
-enumBodyToStage3 :: EnumBody Stage2 I Annot -> Compile Stage3State (EnumBody Stage3 I Annot)
-enumBodyToStage3 (EnumBody constants a) =
-  EnumBody <$> mapM (mapDirectedM enumConstantDeclToStage3) constants <*> pure a
-
-enumConstantDeclToStage3 :: EnumConstantDecl Stage2 I Annot -> Compile Stage3State (EnumConstantDecl Stage3 I Annot)
-enumConstantDeclToStage3 = \case
-  EnumConstantDecl ident expr a -> return $ EnumConstantDecl ident (expressionToStage3 expr) a
-  EnumConstantReserved expr a -> return $ EnumConstantReserved (expressionToStage3 expr) a
-
 packageBodyToStage3 :: PackageBody Stage2 I Annot -> Compile Stage3State (PackageBody Stage3 I Annot)
 packageBodyToStage3 (PackageBody decls a) =
   PackageBody <$> mapM (mapDirectedM fiddleDeclToStage3) decls <*> pure a