feat: Add AdvanceStage typeclass and refactor code to use it

Introduced the `AdvanceStage` typeclass, which provides a mechanism to
transition AST elements between different compilation stages. This
abstraction facilitates easier traversal and modification of the syntax
tree as it progresses through various compilation phases.

2 files changed, 477 insertions, 522 deletions

diff --git a/src/Language/Fiddle/Compiler/Stage1.hs b/src/Language/Fiddle/Compiler/Stage1.hs
index d2fe885..aae80e4 100644
--- a/src/Language/Fiddle/Compiler/Stage1.hs
+++ b/src/Language/Fiddle/Compiler/Stage1.hs
@@ -1,6 +1,7 @@
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE IncoherentInstances #-}
-{-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE UndecidableInstances #-}
 
 module Language.Fiddle.Compiler.Stage1 (toStage2) where
@@ -14,216 +15,133 @@ import Data.Text (Text)
 import qualified Data.Text as Text
 import Data.Type.Bool
 import Debug.Trace
-import GHC.Generics
 import GHC.TypeLits
 import Language.Fiddle.Ast
 import Language.Fiddle.Compiler
 import Language.Fiddle.Types
 import Text.Printf (printf)
 
-newtype Linkage = Linkage Text deriving (Show)
+type Annot = Commented SourceSpan
 
 newtype Path = Path [PathExpression]
 
 newtype PathExpression = PathExpression String
 
-type Annot = Commented SourceSpan
+type M = Compile State
 
 joinPath :: Path -> String
 joinPath (Path l) = intercalate "#" $ reverse (map (\(PathExpression s) -> s) l)
 
+toStage2 :: FiddleUnit Stage1 I Annot -> Compile () (FiddleUnit Stage2 I Annot)
+toStage2 = fmap snd . subCompile (State [] []) . advanceStage (Path mempty)
+
 -- Shorthand for Identity
 type I = Identity
 
-data Stage2CompilerState a
-  = Stage2CompilerState
+newtype Linkage = Linkage Text deriving (Show)
+
+data State
+  = State
       -- Anonymous object type bodies that need to be re-linked
-      ![(Linkage, ObjTypeBody Stage2 I a)]
+      ![(Linkage, ObjTypeBody Stage2 I Annot)]
       -- Anonymous enum bodies that need to be re-linked
-      ![(Linkage, AnonymousBitsType Stage2 I a)]
-
-type M a = Compile (Stage2CompilerState a)
-
-internObjType :: Path -> ObjTypeBody Stage2 I a -> M a (Identifier I a)
-internObjType path body =
-  let str = Text.pack $ joinPath path
-   in do
-        modify $ \(Stage2CompilerState objTypeBodies a) ->
-          Stage2CompilerState ((Linkage str, body) : objTypeBodies) a
-        return (Identifier str (annot body))
-
-internAnonymousBitsType :: Path -> AnonymousBitsType Stage2 I a -> M a (Identifier I a)
-internAnonymousBitsType path anonymousBitsType =
-  let str = Text.pack $ joinPath path
-   in do
-        modify $ \(Stage2CompilerState a anonymousBitsTypes) ->
-          Stage2CompilerState a ((Linkage str, anonymousBitsType) : anonymousBitsTypes)
-        return (Identifier str (annot anonymousBitsType))
-
--- The second stage is a simplified version of the AST without anonymous
--- declarations.
-toStage2 :: FiddleUnit Stage1 I Annot -> Compile () (FiddleUnit Stage2 I Annot)
-toStage2 (FiddleUnit decls annot) = do
-  (s, a) <-
-    subCompile (Stage2CompilerState [] []) $
-      FiddleUnit <$> reconfigureFiddleDecls (Path []) decls <*> pure annot
-  return a
-
-reconfigureFiddleDecls :: Path -> [Directed FiddleDecl Stage1 I Annot] -> M Annot [Directed FiddleDecl Stage2 I Annot]
-reconfigureFiddleDecls p decls = do
-  -- Stage2CompilerState anonymousObjTypes anonymousBitsTypes, decls <- pushState $ do
-  --   put (Stage2CompilerState [] [])
-  --   gets (,) <*> mapM (fiddleDeclToStage2 p) decls
-
-  lastState <- get
-  put (Stage2CompilerState [] [])
-  decls <- mapM (mapDirectedM $ fiddleDeclToStage2 p) decls
-  (Stage2CompilerState anonymousObjTypes anonymousBitsTypes) <- get
-  put lastState
-
-  return $
-    map (asDirected . resolveAnonymousObjType) anonymousObjTypes
-      ++ map (asDirected . resolveAnonymousBitsType) anonymousBitsTypes
-      ++ decls
-  where
-    resolveAnonymousObjType (Linkage linkage, objTypeBody) =
-      ObjTypeDecl (Identifier linkage (annot objTypeBody)) (pure objTypeBody) (annot objTypeBody)
-
-    resolveAnonymousBitsType (Linkage linkage, AnonymousEnumBody expr body a) =
-      BitsDecl (Identifier linkage a) (EnumBitType expr body a) a
-
-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
-
-fiddleDeclToStage2 :: Path -> FiddleDecl Stage1 I Annot -> M Annot (FiddleDecl Stage2 I Annot)
-fiddleDeclToStage2 path decl = do
-  case decl of
-    (OptionDecl i1 i2 a) -> return $ OptionDecl i1 i2 a
-    (PackageDecl n (Identity body) a) -> do
-      (PackageDecl n . Identity <$> packageBodyToStage2 (pushName n path) body) <*> pure a
-    (UsingDecl n a) -> return $ UsingDecl n a
-    (LocationDecl i expr a) -> LocationDecl i <$> toStage2Expr expr <*> pure a
-    (BitsDecl i typ a) -> BitsDecl i <$> bitsTypeToStage2 (pushId i path) typ <*> pure a
-    (ObjTypeDecl i body a) -> ObjTypeDecl i <$> mapM (objTypeBodyToStage2 (pushId i path)) body <*> pure a
-    (ImportDecl importStatement a) -> return $ ImportDecl importStatement a
-    (ObjectDecl i expr typ a) ->
-      ObjectDecl i <$> toStage2Expr expr <*> objectTypeToStage2 (pushId i path) typ <*> pure a
-
-bitsTypeToStage2 :: Path -> BitType Stage1 I Annot -> M Annot (BitType Stage2 I Annot)
-bitsTypeToStage2 path = \case
-  RawBits expr a -> RawBits <$> toStage2Expr expr <*> pure a
-  EnumBitType expr enumBody a ->
-    EnumBitType <$> toStage2Expr expr <*> mapM (enumBodyToStage2 path) enumBody <*> pure a
-
-enumBodyToStage2 :: Path -> EnumBody Stage1 I Annot -> M Annot (EnumBody Stage2 I Annot)
-enumBodyToStage2 path = \case
-  EnumBody constants a -> EnumBody <$> mapM (mapDirectedM (enumConstantToStage2 path)) constants <*> pure a
-
-enumConstantToStage2 :: Path -> EnumConstantDecl Stage1 I Annot -> M Annot (EnumConstantDecl Stage2 I Annot)
-enumConstantToStage2 path = \case
-  EnumConstantDecl i e a -> EnumConstantDecl i <$> toStage2Expr e <*> pure a
-  EnumConstantReserved e a -> EnumConstantReserved <$> toStage2Expr e <*> pure a
-
-objTypeBodyToStage2 :: Path -> ObjTypeBody Stage1 I Annot -> M Annot (ObjTypeBody Stage2 I Annot)
-objTypeBodyToStage2 path (ObjTypeBody bodyType decls annot) =
-  ObjTypeBody bodyType <$> mapM (mapDirectedM $ objTypeDeclToStage2 path) decls <*> pure annot
-
-objTypeDeclToStage2 :: Path -> ObjTypeDecl Stage1 I Annot -> M Annot (ObjTypeDecl Stage2 I Annot)
-objTypeDeclToStage2 path = \case
-  (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
-          <$> objTypeBodyToStage2 path' deferredBody
-          <*> pure maybeIdent
-          <*> pure annot
-  (ReservedDecl expr a) ->
-    ReservedDecl <$> toStage2Expr expr <*> pure a
-  (RegisterDecl maybeModifier maybeIdentifier expression maybeBody annot) ->
-    let path' = maybe path (`pushId` path) maybeIdentifier
-     in RegisterDecl
-          maybeModifier
-          maybeIdentifier
-          <$> toStage2Expr expression
-          <*> mapM (registerBodyToStage2 path') maybeBody
-          <*> pure annot
-
-registerBodyToStage2 :: Path -> RegisterBody Stage1 I Annot -> M Annot (RegisterBody Stage2 I Annot)
-registerBodyToStage2 path (RegisterBody bodyType (Identity (DeferredRegisterBody registerBitsDecl a1)) a2) =
-  RegisterBody bodyType . Identity
-    <$> ( DeferredRegisterBody
-            <$> mapM (mapDirectedM $ registerBitsDeclToStage2 path) registerBitsDecl
-            <*> pure a1
-        )
-    <*> pure a2
-
-registerBitsDeclToStage2 :: Path -> RegisterBitsDecl Stage1 I Annot -> M Annot (RegisterBitsDecl Stage2 I Annot)
-registerBitsDeclToStage2 path = \case
-  ReservedBits expr a -> ReservedBits <$> toStage2Expr expr <*> pure a
-  BitsSubStructure registerBody maybeIdent annot ->
-    let path' = maybe path (`pushId` path) maybeIdent
-     in BitsSubStructure <$> registerBodyToStage2 path' registerBody <*> pure maybeIdent <*> pure annot
-  DefinedBits maybeModifier identifier registerBitsTyperef annot ->
-    let path' = pushId identifier path
-     in ( DefinedBits
-            maybeModifier
-            identifier
-            <$> registerBitsTypeRefToStage2 path' registerBitsTyperef
-            <*> pure annot
-        )
-
-registerBitsTypeRefToStage2 :: Path -> RegisterBitsTypeRef Stage1 I Annot -> M Annot (RegisterBitsTypeRef Stage2 I Annot)
-registerBitsTypeRefToStage2 path = \case
-  RegisterBitsArray typeref expr annot ->
-    RegisterBitsArray
-      <$> registerBitsTypeRefToStage2 path typeref
-      <*> toStage2Expr expr
-      <*> pure annot
-  RegisterBitsReference name annot -> return (RegisterBitsReference name annot)
-  RegisterBitsJustBits expr annot -> RegisterBitsJustBits <$> toStage2Expr expr <*> pure annot
-  RegisterBitsAnonymousType _ anonType annot -> do
-    ident <- internAnonymousBitsType path =<< anonymousBitsTypeToStage2 path anonType
-    return $ RegisterBitsReference (identToName ident) annot
-
-identToName :: Identifier I a -> Name I a
-identToName ident = Name (NonEmpty.singleton ident) (annot ident)
-
-anonymousBitsTypeToStage2 :: Path -> AnonymousBitsType Stage1 I Annot -> M Annot (AnonymousBitsType Stage2 I Annot)
-anonymousBitsTypeToStage2 path = \case
-  AnonymousEnumBody expr (Identity body) annot ->
-    AnonymousEnumBody
-      <$> toStage2Expr expr
-      <*> (Identity <$> enumBodyToStage2 path body)
-      <*> pure annot
-
-objectTypeToStage2 :: Path -> ObjType Stage1 I Annot -> M Annot (ObjType Stage2 I Annot)
-objectTypeToStage2 path = \case
-  (AnonymousObjType _ (Identity body) annot) -> do
-    body' <- objTypeBodyToStage2 path body
-    identifier <- internObjType path body'
-    return (ReferencedObjType (identToName identifier) annot)
-  (ReferencedObjType name annot) -> return $ ReferencedObjType name annot
-  (ArrayObjType objType expr a) ->
-    ArrayObjType <$> objectTypeToStage2 path objType <*> toStage2Expr expr <*> pure a
-
-packageBodyToStage2 :: Path -> PackageBody Stage1 I Annot -> M Annot (PackageBody Stage2 I Annot)
-packageBodyToStage2 p (PackageBody decls a) =
-  PackageBody <$> reconfigureFiddleDecls p decls <*> pure a
-
-toStage2Expr :: Expression Stage1 I Annot -> M Annot (Expression Stage2 I Annot)
-toStage2Expr = \case
-  (Var i a) -> return $ Var i a
-  (LitNum t a) -> LitNum <$> parseNum (unCommented a) t <*> pure a
-
-parseNum :: SourceSpan -> Text -> M a Integer
+      ![(Linkage, AnonymousBitsType Stage2 I Annot)]
+
+instance CompilationStage Stage1 where
+  type StageAfter Stage1 = Stage2
+  type StageMonad Stage1 = M
+  type StageState Stage1 = Path
+  type StageFunctor Stage1 = Identity
+  type StageAnnotation Stage1 = Annot
+
+deriving instance AdvanceStage Stage1 ObjTypeBody
+
+deriving instance AdvanceStage Stage1 DeferredRegisterBody
+
+deriving instance AdvanceStage Stage1 RegisterBody
+
+deriving instance AdvanceStage Stage1 AnonymousBitsType
+
+deriving instance AdvanceStage Stage1 BitType
+
+deriving instance AdvanceStage Stage1 EnumBody
+
+deriving instance AdvanceStage Stage1 EnumConstantDecl
+
+deriving instance (AdvanceStage Stage1 t) => AdvanceStage Stage1 (Directed t)
+
+instance AdvanceStage Stage1 RegisterBitsDecl where
+  modifyState t =
+    return
+      . case t of
+        DefinedBits {definedBitsIdent = i} -> pushId i
+        _ -> id
+
+instance AdvanceStage Stage1 PackageBody where
+  advanceStage p (PackageBody decls a) =
+    PackageBody <$> reconfigureFiddleDecls p decls <*> pure a
+
+instance AdvanceStage Stage1 ObjTypeDecl where
+  modifyState t =
+    return
+      . case t of
+        TypeSubStructure {subStructureName = (Just n)} -> pushId n
+        RegisterDecl {regIdent = (Just n)} -> pushId n
+        _ -> id
+
+instance AdvanceStage Stage1 FiddleDecl where
+  modifyState t =
+    return
+      . case t of
+        PackageDecl {packageName = n} -> pushName n
+        BitsDecl {bitsIdent = i} -> pushId i
+        ObjTypeDecl {objTypeIdent = i} -> pushId i
+        ObjectDecl {objectIdent = i} -> pushId i
+        _ -> id
+
+instance AdvanceStage Stage1 FiddleUnit where
+  advanceStage path (FiddleUnit decls a) =
+    FiddleUnit <$> reconfigureFiddleDecls path decls <*> pure a
+
+instance AdvanceStage Stage1 Expression where
+  advanceStage _ = \case
+    (Var i a) -> return $ Var i a
+    (LitNum t a) -> LitNum <$> parseNum (unCommented a) t <*> pure a
+
+instance AdvanceStage Stage1 RegisterBitsTypeRef where
+  advanceStage path = \case
+    RegisterBitsArray typeref expr annot ->
+      RegisterBitsArray
+        <$> advanceStage path typeref
+        <*> advanceStage path expr
+        <*> pure annot
+    RegisterBitsReference name annot ->
+      return $ RegisterBitsReference name annot
+    RegisterBitsJustBits expr annot ->
+      RegisterBitsJustBits
+        <$> advanceStage path expr
+        <*> pure annot
+    RegisterBitsAnonymousType _ anonType annot -> do
+      ident <-
+        internAnonymousBitsType path
+          =<< advanceStage path anonType
+      return $ RegisterBitsReference (identToName ident) annot
+
+instance AdvanceStage Stage1 ObjType where
+  advanceStage path = \case
+    (AnonymousObjType _ (Identity body) annot) -> do
+      body' <- advanceStage path body
+      identifier <- internObjType path body'
+      return (ReferencedObjType (identToName identifier) annot)
+    (ReferencedObjType name annot) ->
+      return $ ReferencedObjType name annot
+    (ArrayObjType objType expr a) ->
+      ArrayObjType
+        <$> advanceStage path objType
+        <*> advanceStage path expr
+        <*> pure a
+
+parseNum :: SourceSpan -> Text -> Compile s Integer
 parseNum span txt = fromMayberOrFail span "Unable to parse number" $
   case Text.unpack (Text.take 2 txt) of
     "0b" -> toNumWithRadix (Text.drop 2 txt) 2
@@ -255,3 +173,58 @@ parseNum span txt = fromMayberOrFail span "Unable to parse number" $
                       then Nothing
                       else Just (fromIntegral a')
                 )
+
+reconfigureFiddleDecls ::
+  Path ->
+  [Directed FiddleDecl Stage1 I Annot] ->
+  M [Directed FiddleDecl Stage2 I Annot]
+reconfigureFiddleDecls p decls = do
+  lastState <- get
+  put (State [] [])
+  decls <- mapM (mapDirectedM $ advanceStage p) decls
+  (State anonymousObjTypes anonymousBitsTypes) <- get
+  put lastState
+
+  return $
+    map (asDirected . resolveAnonymousObjType) anonymousObjTypes
+      ++ map (asDirected . resolveAnonymousBitsType) anonymousBitsTypes
+      ++ decls
+  where
+    resolveAnonymousObjType (Linkage linkage, objTypeBody) =
+      ObjTypeDecl
+        (Identifier linkage (annot objTypeBody))
+        (pure objTypeBody)
+        (annot objTypeBody)
+
+    resolveAnonymousBitsType (Linkage linkage, AnonymousEnumBody expr body a) =
+      BitsDecl (Identifier linkage a) (EnumBitType expr body a) a
+
+identToName :: Identifier I a -> Name I a
+identToName ident = Name (NonEmpty.singleton ident) (annot ident)
+
+internObjType :: Path -> ObjTypeBody Stage2 I Annot -> M (Identifier I Annot)
+internObjType path body =
+  let str = Text.pack $ joinPath path
+   in do
+        modify $ \(State objTypeBodies a) ->
+          State ((Linkage str, body) : objTypeBodies) a
+        return (Identifier str (annot body))
+
+internAnonymousBitsType ::
+  Path ->
+  AnonymousBitsType Stage2 I Annot ->
+  M (Identifier I Annot)
+internAnonymousBitsType path anonymousBitsType =
+  let str = Text.pack $ joinPath path
+   in do
+        modify $ \(State a anonymousBitsTypes) ->
+          State a ((Linkage str, anonymousBitsType) : anonymousBitsTypes)
+        return (Identifier str (annot anonymousBitsType))
+
+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
diff --git a/src/Language/Fiddle/Compiler/Stage2.hs b/src/Language/Fiddle/Compiler/Stage2.hs
index 2035e3d..1363620 100644
--- a/src/Language/Fiddle/Compiler/Stage2.hs
+++ b/src/Language/Fiddle/Compiler/Stage2.hs
@@ -1,13 +1,14 @@
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE IncoherentInstances #-}
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE UndecidableInstances #-}
 
--- 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 (toStage3) where
 
 import Control.Monad (forM, forM_, unless, when)
+import Control.Monad.Identity (Identity (Identity))
 import Control.Monad.RWS (MonadState (get), MonadWriter (tell), gets, modify')
 import Data.Foldable (Foldable (toList), foldlM)
 import Data.Functor.Identity
@@ -21,13 +22,20 @@ import qualified Data.Map as Map
 import Data.Maybe (fromMaybe)
 import qualified Data.Set as Set
 import qualified Data.Text as Text
-import Data.Word
+import Data.Word (Word32)
 import Language.Fiddle.Ast
 import Language.Fiddle.Compiler
+import Language.Fiddle.Internal.Scopes
 import Language.Fiddle.Types (Commented (unCommented), SourceSpan)
 import Text.Printf (printf)
 import Prelude hiding (unzip)
 
+newtype GlobalState = GlobalState
+  { globalScope :: Scope String (Either SizeBits SizeBytes)
+  }
+
+newtype LocalState = LocalState (ScopePath String)
+
 type I = Identity
 
 type Annot = Commented SourceSpan
@@ -36,344 +44,77 @@ type SizeBits = Word32
 
 type SizeBytes = Word32
 
-data Scope t
-  = Scope
-  { subScopes :: Map String (Scope t),
-    scopeValues :: Map String t
-  }
-
-instance Semigroup (Scope t) where
-  (Scope a1 b1) <> (Scope a2 b2) = Scope (a1 <> a2) (b1 <> b2)
-
-instance Monoid (Scope t) where
-  mempty = Scope mempty mempty
-
-data ScopePath = ScopePath
-  { currentScope :: [String],
-    usingPaths :: [[String]]
-  }
-
-instance Semigroup ScopePath where
-  (ScopePath a1 b1) <> (ScopePath a2 b2) = ScopePath (a1 <> a2) (b1 <> b2)
-
-instance Monoid ScopePath where
-  mempty = ScopePath mempty mempty
-
-emptyScope :: Scope t
-emptyScope = Scope mempty mempty
-
-insertScope :: NonEmpty String -> t -> Scope t -> Scope t
-insertScope (s :| []) v (Scope ss sv) = Scope ss (Map.insert s v sv)
-insertScope (s :| (a : as)) v (Scope ss sv) =
-  Scope
-    ( Map.alter
-        ( \case
-            (fromMaybe emptyScope -> mp) -> Just (insertScope (a :| as) v mp)
-        )
-        s
-        ss
-    )
-    sv
-
-lookupScope :: NonEmpty String -> Scope t -> Maybe t
-lookupScope (s :| []) (Scope _ sv) = Map.lookup s sv
-lookupScope (s :| (a : as)) (Scope ss _) = do
-  subscope <- Map.lookup s ss
-  lookupScope (a :| as) subscope
-
-lookupScopeWithPath :: ScopePath -> NonEmpty String -> Scope t -> Maybe t
-lookupScopeWithPath (ScopePath current others) key scope =
-  let all = reverse (inits current) ++ others
-      e = forM all $ \prefix -> do
-        case lookupScope (NonEmpty.prependList prefix key) scope of
-          Just s -> Left s
-          Nothing -> Right ()
-   in case e of
-        Left v -> Just v
-        Right _ -> Nothing
-
-data Stage3State = Stage3State
-  { inScope :: Scope (Either SizeBits SizeBytes),
-    scopePath :: ScopePath
-  }
-
-insertTypeSize :: Identifier f a -> SizeBits -> Compile Stage3State ()
-insertTypeSize (Identifier s _) size = do
-  modify' $
-    \stage3State ->
-      let fullName =
-            NonEmpty.prependList
-              ((currentScope . scopePath) stage3State)
-              (NonEmpty.singleton (Text.unpack s))
-       in stage3State
-            { inScope =
-                insertScope fullName (Right size) (inScope stage3State)
-            }
-
-lookupTypeSize :: Name I Annot -> Compile Stage3State SizeBits
-lookupTypeSize (Name idents a) = do
-  let path = fmap (\(Identifier s _) -> Text.unpack s) idents
-  scopePath <- gets scopePath
-  mSize <- gets $ lookupScopeWithPath scopePath path . inScope
-  case mSize of
-    Just (Right sz) -> return sz
-    _ -> do
-      tell
-        [ Diagnostic
-            Error
-            ( printf
-                "Cannot resolve %s"
-                (intercalate "." $ NonEmpty.toList path)
-            )
-            (unCommented a)
-        ]
-      compilationFailure
-
-emptyState = Stage3State mempty mempty
-
 toStage3 :: FiddleUnit Stage2 I Annot -> Compile () (FiddleUnit Stage3 I Annot)
-toStage3 (FiddleUnit decls a) =
-  snd
-    <$> subCompile
-      emptyState
-      ( FiddleUnit <$> mapM (mapDirectedM fiddleDeclToStage3) decls <*> pure a
-      )
+toStage3 = fmap snd . subCompile (GlobalState mempty) . advanceStage (LocalState mempty)
+
+instance CompilationStage Stage2 where
+  type StageAfter Stage2 = Stage3
+  type StageMonad Stage2 = Compile GlobalState
+  type StageState Stage2 = LocalState
+  type StageFunctor Stage2 = Identity
+  type StageAnnotation Stage2 = Commented SourceSpan
+
+deriving instance AdvanceStage Stage2 FiddleUnit
+
+deriving instance AdvanceStage Stage2 Expression
+
+instance AdvanceStage Stage2 FiddleDecl where
+  modifyState t s = case t of
+    (BitsDecl id typ a) -> do
+      typeSize <- getTypeSize typ
+      insertTypeSize s id typeSize
+      return s
+    (PackageDecl n _ _) -> do
+      let strs = nameToList n
+      let (LocalState scopePath) = s
+
+      return $
+        LocalState $
+          scopePath {currentScope = strs ++ currentScope scopePath}
+    (UsingDecl n _) ->
+      let (LocalState scopePath) = s
+       in return $
+            LocalState $
+              scopePath
+                { usingPaths = nameToList n : usingPaths scopePath
+                }
+    _ -> return s
 
-exprToSize ::
-  (NumberType stage ~ Integer) =>
-  Expression stage I Annot ->
-  Compile s Integer
-exprToSize (LitNum num _) = return num
-exprToSize e = do
-  tell [Diagnostic Error "Variables not allowed" (unCommented $ annot e)]
-  compilationFailure
+nameToList :: Name f a -> [String]
+nameToList (Name idents _) = map (\(Identifier (Text.unpack -> s) _) -> s) (toList idents)
 
-getTypeSize :: BitType Stage2 I Annot -> Compile s SizeBits
-getTypeSize (RawBits expr _) = fromIntegral <$> exprToSize expr
-getTypeSize (EnumBitType expr (Identity (EnumBody constants _)) ann) = do
-  declaredSize <- fromIntegral <$> exprToSize expr
+instance AdvanceStage Stage2 ObjTypeBody where
+  advanceStage s body = fst <$> objTypeBodyToStage3 s body 0
 
-  -- If the declared size is less than or equal to 4, we'll enforce that the
-  -- enum is packed. This is to make sure the user has covered all bases.
-  when (declaredSize <= 4) $ do
-    imap <-
-      foldlM
-        ( \imap (undirected -> enumConst) -> do
-            number <- case enumConst of
-              EnumConstantDecl _ expr _ -> exprToSize expr
-              EnumConstantReserved expr _ -> exprToSize expr
+deriving instance AdvanceStage Stage2 ObjType
 
-            when (number >= 2 ^ declaredSize) $
-              tell
-                [ Diagnostic
-                    Error
-                    ( printf
-                        "Enum constant too large. Max allowed %d\n"
-                        ((2 :: Int) ^ declaredSize)
-                    )
-                    (unCommented (annot enumConst))
-                ]
+deriving instance AdvanceStage Stage2 DeferredRegisterBody
 
-            return $ IntMap.insert (fromIntegral number) True imap
-        )
-        IntMap.empty
-        constants
-    let missing =
-          filter (not . (`IntMap.member` imap)) [0 .. 2 ^ declaredSize - 1]
-    unless (null missing) $
-      tell
-        [ Diagnostic
-            Warning
-            ( printf
-                "Missing enum constants %s. Small enums should be fully \
-                \ populated. Use 'reserved' if needed."
-                (intercalate ", " (map show missing))
-            )
-            (unCommented ann)
-        ]
+deriving instance AdvanceStage Stage2 RegisterBitsDecl
 
-  return declaredSize
+instance AdvanceStage Stage2 RegisterBody where
+  advanceStage s body = fst <$> registerBodyToStage3 s body
 
-addCurrentScope :: [String] -> Compile Stage3State ()
-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 = \case
-  OptionDecl i1 i2 a -> return $ OptionDecl i1 i2 a
-  PackageDecl n@(Name idents _) body a -> do
-    let strs = map (\(Identifier (Text.unpack -> s) _) -> s) (toList idents)
-    Stage3State {scopePath = savedScopePath} <- get
-    addCurrentScope strs
-    PackageDecl n
-      <$> mapM packageBodyToStage3 body
-      <*> pure a
-      <* modify' (\st -> st {scopePath = savedScopePath})
-  UsingDecl n@(Name idents _) a -> do
-    let strs = map (\(Identifier t _) -> Text.unpack t) (toList idents)
-    modify'
-      ( \st ->
-          let (ScopePath cur using) = scopePath st
-           in st
-                { scopePath = ScopePath cur (strs : using)
-                }
-      )
-    return $ UsingDecl n a
-  LocationDecl id expr a -> return $ LocationDecl id (switchStage expr) a
-  BitsDecl id typ a -> do
-    typeSize <- getTypeSize typ
-    insertTypeSize id typeSize
-    return $ BitsDecl id (switchStage typ) a
-  ObjTypeDecl ident body 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
-      (switchStage expr)
-      <$> objTypeToStage3 typ
-      <*> pure a
-
-objTypeToStage3 ::
-  ObjType Stage2 I Annot -> Compile Stage3State (ObjType Stage3 I Annot)
-objTypeToStage3 = \case
-  ArrayObjType objtype expr a ->
-    ArrayObjType
-      <$> objTypeToStage3 objtype
-      <*> pure (switchStage expr)
-      <*> pure a
-  ReferencedObjType ident a -> return $ ReferencedObjType ident a
+deriving instance AdvanceStage Stage2 RegisterBitsTypeRef
 
-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
+deriving instance AdvanceStage Stage2 AnonymousBitsType
 
-    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
+deriving instance AdvanceStage Stage2 BitType
 
-                    if isUnion
-                      then checkUnion cursor size (s3 : returned) a
-                      else
-                        return (cursor + size, s3 : returned)
-            )
-            (0, [])
-            decls
+deriving instance AdvanceStage Stage2 EnumBody
 
-        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 ()
+deriving instance AdvanceStage Stage2 EnumConstantDecl
 
-checkUnion :: Word32 -> Word32 -> b -> Commented SourceSpan -> Compile Stage3State (Word32, b)
-checkUnion cursor subsize ret a = do
-  when (cursor /= 0 && subsize /= cursor) $ do
-    tell
-      [ Diagnostic
-          Warning
-          ( printf
-              "Jagged union found. Found size %d, expected %d.\n \
-              \ Please wrap smaller fields in a struct with padding so all \
-              \ fields are the same size?"
-              subsize
-              cursor
-          )
-          (unCommented a)
-      ]
-  return (max cursor subsize, ret)
+deriving instance AdvanceStage Stage2 PackageBody
 
-registerBitsTypeRefToStage3 ::
-  RegisterBitsTypeRef Stage2 I Annot ->
-  Compile Stage3State (RegisterBitsTypeRef Stage3 I Annot, Word32)
-registerBitsTypeRefToStage3 = \case
-  RegisterBitsArray ref expr a -> do
-    (ref', size) <- registerBitsTypeRefToStage3 ref
-    multiplier <- exprToSize expr
-    return
-      ( RegisterBitsArray ref' (switchStage expr) a,
-        size * fromIntegral multiplier
-      )
-  RegisterBitsReference name a ->
-    (RegisterBitsReference name a,) <$> lookupTypeSize name
-  RegisterBitsJustBits expr a ->
-    (RegisterBitsJustBits (switchStage expr) a,)
-      . fromIntegral
-      <$> exprToSize expr
+deriving instance (AdvanceStage Stage2 t) => AdvanceStage Stage2 (Directed t)
 
 objTypeBodyToStage3 ::
-  ObjTypeBody Stage2 I Annot -> Word32 -> Compile Stage3State (ObjTypeBody Stage3 I Annot, Word32)
-objTypeBodyToStage3 (ObjTypeBody bodyType decls a) startOff = do
+  LocalState ->
+  ObjTypeBody Stage2 I Annot ->
+  Word32 ->
+  Compile GlobalState (ObjTypeBody Stage3 I Annot, Word32)
+objTypeBodyToStage3 st (ObjTypeBody bodyType decls a) startOff = do
   let isUnion = case bodyType of
         Union {} -> True
         _ -> False
@@ -383,7 +124,9 @@ objTypeBodyToStage3 (ObjTypeBody bodyType decls a) startOff = do
           case undirected decl of
             RegisterDecl mMod mIdent expr mBody a -> do
               (s3RegisterBody, mCalculatedSize) <-
-                fUnzip <$> mapM registerBodyToStage3 mBody
+                fUnzip <$> mapM (registerBodyToStage3 st) mBody
+
+              nExpr <- advanceStage st expr
 
               let s3 =
                     mapDirected
@@ -391,7 +134,7 @@ objTypeBodyToStage3 (ObjTypeBody bodyType decls a) startOff = do
                           RegisterDecl
                             mMod
                             mIdent
-                            (switchStage expr)
+                            nExpr
                             s3RegisterBody
                             a
                       )
@@ -436,6 +179,7 @@ objTypeBodyToStage3 (ObjTypeBody bodyType decls a) startOff = do
             TypeSubStructure (Identity subBody) maybeIdent annot -> do
               (newBody, size) <-
                 objTypeBodyToStage3
+                  st
                   subBody
                   ( if isUnion then startOff else cursor
                   )
@@ -456,8 +200,10 @@ objTypeBodyToStage3 (ObjTypeBody bodyType decls a) startOff = do
                       "Can only reserve a multiple of 8 bits in this context."
                       (unCommented a)
                   ]
+
+              expr' <- advanceStage st expr
               let size = size' `div` 8
-              let s3 = mapDirected (const $ ReservedDecl (switchStage expr) annot) decl
+              let s3 = mapDirected (const $ ReservedDecl expr' annot) decl
               if isUnion
                 then
                   checkUnion cursor size (s3 : returned) a
@@ -517,6 +263,242 @@ objTypeBodyToStage3 (ObjTypeBody bodyType decls a) startOff = do
     pushApply (Just (a, b)) = (Just a, Just b)
     pushApply Nothing = (Nothing, Nothing)
 
-packageBodyToStage3 :: PackageBody Stage2 I Annot -> Compile Stage3State (PackageBody Stage3 I Annot)
-packageBodyToStage3 (PackageBody decls a) =
-  PackageBody <$> mapM (mapDirectedM fiddleDeclToStage3) decls <*> pure a
+registerBodyToStage3 ::
+  LocalState ->
+  RegisterBody Stage2 I Annot ->
+  Compile GlobalState (RegisterBody Stage3 I Annot, Word32)
+registerBodyToStage3
+  st
+  (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
+                    expr' <- advanceStage st expr
+                    let s3 =
+                          mapDirected
+                            (const $ ReservedBits 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 st 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 st 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
+          )
+    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 ()
+
+registerBitsTypeRefToStage3 ::
+  LocalState ->
+  RegisterBitsTypeRef Stage2 I Annot ->
+  Compile GlobalState (RegisterBitsTypeRef Stage3 I Annot, Word32)
+registerBitsTypeRefToStage3 localState = \case
+  RegisterBitsArray ref expr a -> do
+    (ref', size) <- registerBitsTypeRefToStage3 localState ref
+    multiplier <- exprToSize expr
+    expr' <- advanceStage localState expr
+    return
+      ( RegisterBitsArray ref' expr' a,
+        size * fromIntegral multiplier
+      )
+  RegisterBitsReference name a ->
+    (RegisterBitsReference name a,) <$> lookupTypeSize localState name
+  RegisterBitsJustBits expr a -> do
+    expr' <- advanceStage localState expr
+    (RegisterBitsJustBits expr' a,)
+      . fromIntegral
+      <$> exprToSize expr
+
+checkUnion :: Word32 -> Word32 -> b -> Commented SourceSpan -> Compile s (Word32, b)
+checkUnion cursor subsize ret a = do
+  when (cursor /= 0 && subsize /= cursor) $ do
+    tell
+      [ Diagnostic
+          Warning
+          ( printf
+              "Jagged union found. Found size %d, expected %d.\n \
+              \ Please wrap smaller fields in a struct with padding so all \
+              \ fields are the same size?"
+              subsize
+              cursor
+          )
+          (unCommented a)
+      ]
+  return (max cursor subsize, ret)
+
+exprToSize ::
+  (NumberType stage ~ Integer) =>
+  Expression stage I Annot ->
+  Compile s Integer
+exprToSize (LitNum num _) = return num
+exprToSize e = do
+  tell [Diagnostic Error "Variables not allowed" (unCommented $ annot e)]
+  compilationFailure
+
+lookupTypeSize :: LocalState -> Name I Annot -> Compile GlobalState SizeBits
+lookupTypeSize (LocalState scopePath) (Name idents a) = do
+  -- Convert the list of identifiers to a string path
+  let path = fmap (\(Identifier s _) -> Text.unpack s) idents
+
+  -- Get the current scope and perform the lookup
+  results <- gets $ lookupScopeWithPath scopePath path . globalScope
+
+  case results of
+    -- Successfully resolved to a unique size
+    [(_, Right sz)] -> return sz
+    -- Multiple ambiguous results found
+    matches@(_ : _) -> do
+      -- Generate a list of ambiguous paths for error reporting
+      let ambiguousPaths =
+            map
+              ( \(resolvedPath, _) ->
+                  intercalate "." (NonEmpty.toList resolvedPath)
+              )
+              matches
+      tell
+        [ Diagnostic
+            Error
+            ( printf
+                "Ambiguous occurrence of '%s'. Multiple matches found:\n%s"
+                (intercalate "." $ NonEmpty.toList path)
+                (unlines ambiguousPaths) -- List all ambiguous paths
+            )
+            (unCommented a)
+        ]
+      compilationFailure
+
+    -- No matches found
+    _ -> do
+      tell
+        [ Diagnostic
+            Error
+            ( printf
+                "Cannot resolve '%s'. No matching symbols found."
+                (intercalate "." $ NonEmpty.toList path)
+            )
+            (unCommented a)
+        ]
+      compilationFailure
+
+getTypeSize :: BitType Stage2 I Annot -> Compile s SizeBits
+getTypeSize (RawBits expr _) = fromIntegral <$> exprToSize expr
+getTypeSize (EnumBitType expr (Identity (EnumBody constants _)) ann) = do
+  declaredSize <- fromIntegral <$> exprToSize expr
+
+  -- If the declared size is less than or equal to 4, we'll enforce that the
+  -- enum is packed. This is to make sure the user has covered all bases.
+  when (declaredSize <= 4) $ do
+    imap <-
+      foldlM
+        ( \imap (undirected -> enumConst) -> do
+            number <- case enumConst of
+              EnumConstantDecl _ expr _ -> exprToSize expr
+              EnumConstantReserved expr _ -> exprToSize expr
+
+            when (number >= 2 ^ declaredSize) $
+              tell
+                [ Diagnostic
+                    Error
+                    ( printf
+                        "Enum constant too large. Max allowed %d\n"
+                        ((2 :: Int) ^ declaredSize)
+                    )
+                    (unCommented (annot enumConst))
+                ]
+
+            return $ IntMap.insert (fromIntegral number) True imap
+        )
+        IntMap.empty
+        constants
+    let missing =
+          filter (not . (`IntMap.member` imap)) [0 .. 2 ^ declaredSize - 1]
+    unless (null missing) $
+      tell
+        [ Diagnostic
+            Warning
+            ( printf
+                "Missing enum constants %s. Small enums should be fully \
+                \ populated. Use 'reserved' if needed."
+                (intercalate ", " (map show missing))
+            )
+            (unCommented ann)
+        ]
+
+  return declaredSize
+
+insertTypeSize :: LocalState -> Identifier f a -> SizeBits -> Compile GlobalState ()
+insertTypeSize (LocalState scopePath) (Identifier s _) size = do
+  modify' $
+    \(GlobalState globalScope) ->
+      let fullName =
+            NonEmpty.prependList
+              (currentScope scopePath)
+              (NonEmpty.singleton (Text.unpack s))
+       in GlobalState $
+            insertScope fullName (Right size) globalScope