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.

9 files changed, 1334 insertions, 1048 deletions

diff --git a/src/Language/Fiddle/Ast/FileInterface.hs b/src/Language/Fiddle/Ast/FileInterface.hs
index d29fc9d..c1cfac8 100644
--- a/src/Language/Fiddle/Ast/FileInterface.hs
+++ b/src/Language/Fiddle/Ast/FileInterface.hs
@@ -6,28 +6,33 @@ module Language.Fiddle.Ast.FileInterface where
 -- the import statements should supply an fdi (fiddle interface) file to speed
 -- up subsequent compilations.
 
--- import Data.Text
--- 
--- data ObjectType = ObjectType
---   { objectTypeSize :: Word32
---   }
--- 
--- data Metatype
---   = Object
---       { objectLocation :: Word64,
---         objectType :: Text
---       }
---   | Type
---       { typeSizeBytes :: Word32
---       }
--- 
--- data Element a = Element
---   { elementFullyQualifiedSymbol :: Text,
---     elementDocumentation :: Maybe Text,
---     elementMetatype :: Metatype,
---     elementAnnotation :: a
---   }
--- 
--- data FileInterface a = FiddleInterface
---   { exportedElements :: [Element a]
---   }
+import Data.Text
+import Data.Word
+
+data ObjectType = ObjectType
+  { objectTypeSize :: Word32
+  }
+
+data Metatype
+  = Object
+      { objectLocation :: Word64,
+        objectType :: Text
+      }
+  | Type
+      { typeSizeBytes :: Word32
+      }
+
+data Element a = Element
+  { elementFullyQualifiedSymbol :: Text,
+    elementDocumentation :: Maybe Text,
+    elementMetatype :: Metatype,
+    elementAnnotation :: a
+  }
+
+data ResolvedImport a = ResolvedImport {
+  dependencies :: [String]
+}
+
+data FileInterface a = FiddleInterface
+  { exportedElements :: [Element a]
+  }
diff --git a/src/Language/Fiddle/Ast/Internal/Instances.hs b/src/Language/Fiddle/Ast/Internal/Instances.hs
index c8c606c..2f3707e 100644
--- a/src/Language/Fiddle/Ast/Internal/Instances.hs
+++ b/src/Language/Fiddle/Ast/Internal/Instances.hs
@@ -15,27 +15,6 @@ import GHC.TypeError as TypeError
 import Language.Fiddle.Ast.Internal.Kinds
 import Language.Fiddle.Ast.Internal.Stage
 
--- Type class to easily switch between stages if there is no difference in the
--- syntax tree structure between these stages. Can make things much cleaner and
--- avoids the boilerplate and bugprone-ness of needing to rote copy everything.
-class
-  (Functor f) =>
-  EasySwitchStage
-    (t :: Stage -> SynTree)
-    (f :: Type -> Type)
-    (fromStage :: Stage)
-    (toStage :: Stage)
-  where
-  switchStage :: t fromStage f a -> t toStage f a
-  default switchStage ::
-    ( Generic (t fromStage f a),
-      Generic (t toStage f a),
-      GEasySwitchStage (Rep (t fromStage f a)) (Rep (t toStage f a))
-    ) =>
-    t fromStage f a ->
-    t toStage f a
-  switchStage t = to $ gSwitchStage (from t)
-
 -- Class for walking a syntax tree under the context of a monad and modifying
 -- the different parts of the SynTree type..
 class Alter (t :: SynTree) where
@@ -63,6 +42,154 @@ class Alter (t :: SynTree) where
     m (t f2 a2)
   alter ffn fn t = to <$> galter (proxyOf t) ffn fn (from t)
 
+-- | 'CompilationStage' is a type class representing a stage in the compiler
+--   pipeline. Each 'stage' has associated types that define how it transforms
+--   syntax trees, manages state, and handles annotations.
+--
+--   This class requires:
+--   * A 'StageMonad' for handling monadic actions during the compilation process.
+--   * A 'StageFunctor' for traversing the syntax tree.
+--   * A 'StageState' to represent local state information during traversal.
+class
+  (Monad (StageMonad stage), Traversable (StageFunctor stage)) =>
+  CompilationStage stage
+  where
+  -- | The next stage in the compilation pipeline.
+  type StageAfter stage :: Stage
+
+  -- | The monadic context for this stage (e.g., 'Compile' or 'Either').
+  type StageMonad stage :: Type -> Type
+
+  -- | The state type associated with this stage. This state is designed to be
+  --   used as a 'local state', meaning any modifications to this state are
+  --   only visible to child nodes during the traversal. To share state across
+  --   sibling nodes, use the 'StageMonad'.
+  type StageState stage :: Type
+
+  -- | The functor used for traversing the syntax tree during this stage.
+  type StageFunctor stage :: Type -> Type
+
+  -- | The type of annotations associated with nodes in the syntax tree.
+  type StageAnnotation stage :: Type
+
+-- | Utility type that captures the structure of a syntax tree node for a given
+--   stage. This type alias simplifies references to the full tree structure in
+--   other parts of the code.
+type TreeType (t :: StagedSynTree) (s :: Stage) =
+  t s (StageFunctor s) (StageAnnotation s)
+
+-- | 'AdvanceStage' defines how to transform an Abstract Syntax Tree (AST) node
+--   from one stage to the next in the compiler pipeline. This transformation
+--   can be customized per node type, or a default generic implementation can be
+--   used.
+class
+  (CompilationStage stage) =>
+  AdvanceStage (stage :: Stage) (t :: StagedSynTree)
+  where
+  -- | 'advanceStage' transitions a syntax tree node from the current stage to
+  --   the next stage ('StageAfter stage'). It takes the current local state
+  --   ('StageState') and the syntax tree node ('TreeType') as input and returns
+  --   a monadic action that produces the transformed node for the next stage.
+  --
+  --   This function typically performs tree transformations, checks, or other
+  --   modifications necessary for the compilation process.
+  --
+  --   Parameters:
+  --   * 'StageState stage' - The local state for this stage of the compiler.
+  --     This state is only visible to child nodes. Any changes do not affect
+  --     sibling nodes.
+  --   * 'TreeType t stage' - The syntax tree node at the current stage.
+  --
+  --   Returns: A monadic action that produces the transformed syntax tree node
+  --   at the next stage.
+  advanceStage ::
+    StageState stage -> -- Local state for the current stage
+    TreeType t stage -> -- Syntax tree node at the current stage
+    StageMonad
+      stage
+      ( t
+          (StageAfter stage) -- The next stage in the pipeline
+          (StageFunctor stage) -- The functor associated with the next stage
+          (StageAnnotation stage) -- Annotation type for the next stage
+      )
+
+  -- | Default implementation of 'advanceStage' using generics. This leverages
+  --   'GAdvanceStage' to automatically traverse and transform the syntax tree.
+  --   If an AST type derives 'Generic', this default can be used to reduce
+  --   boilerplate code.
+  --
+  --   Before performing the generic transformation with 'gAdvanceStage', the
+  --   'modifyState' function is called to potentially alter the local state
+  --   based on the current node.
+  --
+  --   This implementation is useful for cases where the tree structure remains
+  --   mostly unchanged between stages, and only the state or annotations need
+  --   to be adjusted.
+  default advanceStage ::
+    ( GAdvanceStage
+        (StageState stage) -- The local state for this stage
+        (StageMonad stage) -- The monadic context of this stage
+        (Rep (TreeType t stage)) -- Generic representation of the current tree type
+        (Rep (t (StageAfter stage) (StageFunctor stage) (StageAnnotation stage))), -- Generic representation of the next stage's tree type
+      Generic (TreeType t stage), -- The current tree type must be an instance of 'Generic'
+      Generic
+        ( t
+            (StageAfter stage) -- The tree type at the next stage
+            (StageFunctor stage) -- The functor for the next stage
+            (StageAnnotation stage) -- The annotation type for the next stage
+        )
+    ) =>
+    StageState stage -> -- Local state for the current stage
+    TreeType t stage -> -- Syntax tree node at the current stage
+    StageMonad
+      stage
+      ( t
+          (StageAfter stage) -- The tree type for the next stage
+          (StageFunctor stage) -- The functor for the next stage
+          (StageAnnotation stage) -- The annotation type for the next stage
+      )
+  advanceStage s t = do
+    -- Modify the local state for this node before performing the transformation
+    s' <- modifyState t s
+    -- Perform the generic transformation using 'gAdvanceStage'
+    to <$> gAdvanceStage s' (from t)
+
+  -- | 'modifyState' allows for changes to the local state ('StageState') before
+  --   transforming the syntax tree node. This is called on each node during the
+  --   traversal, allowing the state to be adjusted based on the current node's
+  --   information.
+  --
+  --   The default implementation simply returns the unchanged state, but it can
+  --   be overridden to implement custom state modifications.
+  --
+  --   Note: This state modification is local to the current node and its
+  --   children. Changes to this state are not visible to sibling nodes. If
+  --   information needs to be shared across siblings, consider using the
+  --   'StageMonad' for that purpose.
+  --
+  --   'modifyState' is only called by the default implementation of
+  --   advanceStage. If 'advaceStage' is overridden, then overriding this
+  --   function is superfluous.
+  --
+  --   Parameters:
+  --   * 'TreeType t stage' - The syntax tree node at the current stage.
+  --   * 'StageState stage' - The local state for this stage.
+  --
+  --   Returns: A monadic action that produces the (potentially modified) state.
+  modifyState ::
+    TreeType t stage -> -- Syntax tree node at the current stage
+    StageState stage -> -- Local state for the current stage
+    StageMonad stage (StageState stage) -- The modified local state
+  modifyState _ = return
+
+-- | 'GAdvanceStage' is a helper type class that performs the transformation
+--   of the generic representation of a syntax tree node. It is used by the
+--   default implementation of 'advanceStage' to traverse and modify nodes
+--   automatically.
+class GAdvanceStage s m from to where
+  gAdvanceStage :: s -> from x -> m (to x)
+
+
 -- A syntax tree object is annotated if it has an annotation 'a' as the last
 -- element.
 class Annotated (t :: SynTree) where
@@ -170,59 +297,154 @@ instance (GAnnot a r, GAnnot a l) => GAnnot a (l :+: r) where
 instance (GAnnot a r) => GAnnot a (M1 i c r) where
   gannot (M1 a) = gannot a
 
+-- instance
+--   (GEasySwitchStage s1 s2) =>
+--   GEasySwitchStage (M1 i c s1) (M1 i c s2)
+--   where
+--   gSwitchStage (M1 a) = M1 (gSwitchStage a)
+-- 
+-- instance
+--   (GEasySwitchStage l1 l2, GEasySwitchStage r1 r2) =>
+--   GEasySwitchStage (l1 :+: r1) (l2 :+: r2)
+--   where
+--   gSwitchStage (R1 r) = R1 $ gSwitchStage r
+--   gSwitchStage (L1 l) = L1 $ gSwitchStage l
+-- 
+-- instance
+--   (GEasySwitchStage l1 l2, GEasySwitchStage r1 r2) =>
+--   (GEasySwitchStage (l1 :*: r1) (l2 :*: r2))
+--   where
+--   gSwitchStage (l :*: r) = gSwitchStage l :*: gSwitchStage r
+-- 
+-- instance
+--   (EasySwitchStage t f fs ts) =>
+--   (GEasySwitchStage (Rec0 (t fs f a)) (Rec0 (t ts f a)))
+--   where
+--   gSwitchStage (K1 val) = K1 (switchStage val)
+-- 
+-- instance
+--   ( EasySwitchStage t f fs ts,
+--     Functor func
+--   ) =>
+--   (GEasySwitchStage (Rec0 (func (t fs f a))) (Rec0 (func (t ts f a))))
+--   where
+--   gSwitchStage (K1 val) = K1 (switchStage <$> val)
+-- 
+-- instance (GEasySwitchStage (Rec0 a) (Rec0 a)) where
+--   gSwitchStage = id
+-- 
+-- instance
+--   ( TypeError
+--       ( TypeError.Text "Unable to match type "
+--           :<>: TypeError.ShowType a
+--           :<>: TypeError.Text " with "
+--           :<>: TypeError.ShowType b
+--       )
+--   ) =>
+--   (GEasySwitchStage (Rec0 a) (Rec0 b))
+--   where
+--   gSwitchStage = error "Cannot be called"
+
+-- class GEasySwitchStage r1 r2 where
+--   gSwitchStage :: r1 x -> r2 x
+
+proxyOf :: t f a -> Proxy t
+proxyOf _ = Proxy
+
+instance (Alter t, Traversable f) => Functor (t f) where
+  fmap f t = runIdentity (alter return (return . f) t)
+
+-- | 'GAdvanceStage' instance for metadata wrappers ('M1').
+--   This instance allows advancing the stage of a metadata node in a generic
+--   representation. The metadata node ('M1') wraps another node ('s1'), which
+--   is recursively advanced to the next stage using 'gAdvanceStage'.
 instance
-  (GEasySwitchStage s1 s2) =>
-  GEasySwitchStage (M1 i c s1) (M1 i c s2)
+  (Monad m, GAdvanceStage s m s1 s2) =>
+  GAdvanceStage s m (M1 i c s1) (M1 i c s2)
   where
-  gSwitchStage (M1 a) = M1 (gSwitchStage a)
+  gAdvanceStage s (M1 a) = M1 <$> gAdvanceStage s a
 
+-- | 'GAdvanceStage' instance for sum types (':+:'). This handles the case
+--   where the generic representation of a type is a sum (i.e., an 'Either'-like
+--   choice between two alternatives). The sum type can be either 'L1' (left)
+--   or 'R1' (right), and 'gAdvanceStage' is called recursively on the selected
+--   branch.
 instance
-  (GEasySwitchStage l1 l2, GEasySwitchStage r1 r2) =>
-  GEasySwitchStage (l1 :+: r1) (l2 :+: r2)
+  (Monad m, GAdvanceStage s m l1 l2, GAdvanceStage s m r1 r2) =>
+  GAdvanceStage s m (l1 :+: r1) (l2 :+: r2)
   where
-  gSwitchStage (R1 r) = R1 $ gSwitchStage r
-  gSwitchStage (L1 l) = L1 $ gSwitchStage l
+  gAdvanceStage s (R1 r) = R1 <$> gAdvanceStage s r
+  gAdvanceStage s (L1 l) = L1 <$> gAdvanceStage s l
 
+-- | 'GAdvanceStage' instance for product types (':*:'). This handles cases
+--   where the generic representation of a type is a product (i.e., a tuple of
+--   multiple components). It recursively advances each component ('l' and 'r')
+--   to the next stage.
 instance
-  (GEasySwitchStage l1 l2, GEasySwitchStage r1 r2) =>
-  (GEasySwitchStage (l1 :*: r1) (l2 :*: r2))
+  (Monad m, GAdvanceStage s m l1 l2, GAdvanceStage s m r1 r2) =>
+  GAdvanceStage s m (l1 :*: r1) (l2 :*: r2)
   where
-  gSwitchStage (l :*: r) = gSwitchStage l :*: gSwitchStage r
+  gAdvanceStage s (l :*: r) =
+    (:*:) <$> gAdvanceStage s l <*> gAdvanceStage s r
 
+-- | 'GAdvanceStage' instance for record fields ('Rec0') containing a single
+--   AST element ('t') to be advanced. This instance covers the case where the
+--   field is an individual AST node that implements 'AdvanceStage'.
+--   It advances this node using 'advanceStage'.
 instance
-  (EasySwitchStage t f fs ts) =>
-  (GEasySwitchStage (Rec0 (t fs f a)) (Rec0 (t ts f a)))
+  ( Monad m,
+    AdvanceStage stage t',
+    Traversable f,
+    StageAfter stage ~ stage',
+    StageMonad stage ~ m,
+    StageState stage ~ s,
+    StageFunctor stage ~ f,
+    StageAnnotation stage ~ a
+  ) =>
+  GAdvanceStage s m (Rec0 (t' stage f a)) (Rec0 (t' stage' f a))
   where
-  gSwitchStage (K1 val) = K1 (switchStage val)
+  gAdvanceStage st (K1 val) = K1 <$> advanceStage st val
 
+-- | 'GAdvanceStage' instance for record fields ('Rec0') containing a functor
+--   ('func') of AST elements ('t'). This handles cases where the field is a
+--   container (e.g., list, 'Maybe') of AST nodes that need to be advanced.
+--   Each node in the container is transformed using 'advanceStage'.
 instance
-  ( EasySwitchStage t f fs ts,
-    Functor func
+  ( Monad m,
+    AdvanceStage stage t',
+    Traversable f,
+    StageAfter stage ~ stage',
+    StageMonad stage ~ m,
+    StageState stage ~ s,
+    StageFunctor stage ~ f,
+    StageAnnotation stage ~ a,
+    Traversable func
   ) =>
-  (GEasySwitchStage (Rec0 (func (t fs f a))) (Rec0 (func (t ts f a))))
+  GAdvanceStage s m (Rec0 (func (t' stage f a))) (Rec0 (func (t' stage' f a)))
   where
-  gSwitchStage (K1 val) = K1 (switchStage <$> val)
-
-instance (GEasySwitchStage (Rec0 a) (Rec0 a)) where
-  gSwitchStage = id
+  gAdvanceStage st (K1 val) = K1 <$> mapM (advanceStage st) val
 
+-- | 'GAdvanceStage' instance for record fields ('Rec0') containing a
+--   functor ('f') wrapping an AST element. This handles cases where the field
+--   is a container ('f') of AST nodes that need to be advanced. Each node in
+--   the container is transformed using 'advanceStage'.
 instance
-  ( TypeError
-      ( TypeError.Text "Unable to match type "
-          :<>: TypeError.ShowType a
-          :<>: TypeError.Text " with "
-          :<>: TypeError.ShowType b
-      )
+  ( Monad m,
+    AdvanceStage stage t',
+    Traversable f,
+    StageAfter stage ~ stage',
+    StageMonad stage ~ m,
+    StageState stage ~ s,
+    StageFunctor stage ~ f,
+    StageAnnotation stage ~ a
   ) =>
-  (GEasySwitchStage (Rec0 a) (Rec0 b))
+  GAdvanceStage s m (Rec0 (f (t' stage f a))) (Rec0 (f (t' stage' f a)))
   where
-  gSwitchStage = error "Cannot be called"
-
-class GEasySwitchStage r1 r2 where
-  gSwitchStage :: r1 x -> r2 x
+  gAdvanceStage st (K1 val) = K1 <$> mapM (advanceStage st) val
 
-proxyOf :: t f a -> Proxy t
-proxyOf _ = Proxy
-
-instance (Alter t, Traversable f) => Functor (t f) where
-  fmap f t = runIdentity (alter return (return . f) t)
+-- | 'GAdvanceStage' instance for simple record fields ('Rec0') that do not
+--   need to change between stages. This is used for fields that are not AST
+--   nodes and remain the same when advancing the stage (e.g., primitive
+--   types like 'Int', 'Bool', etc.).
+instance (Monad m) => GAdvanceStage s m (Rec0 a) (Rec0 a) where
+  gAdvanceStage _ (K1 val) = return (K1 val)
diff --git a/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs b/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
index 48852ee..827f712 100644
--- a/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
+++ b/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
@@ -109,429 +109,400 @@ data Name :: SynTree where
   Name :: NonEmpty (Identifier f a) -> a -> Name f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- Syntax tree fo the directive sublanguage. Directives can be on many elements
--- and provide the compiler with additional information about the emitted code.
---
--- The directive subtree by design does not depend on the compilation stage.
--- This is because the directive sublanguage should pass directly to the backend
--- compilation stage.
+-- | Represents a directive in the Fiddle language. A directive provides 
+--   additional metadata or instructions that the compiler can use during 
+--   code generation. Directives can be attached to many elements in the 
+--   syntax tree.
 data Directive :: SynTree where
-  Directive :: f (DirectiveBody f a) -> a -> Directive f a
+  Directive ::
+    { directiveBody :: f (DirectiveBody f a), -- ^ The body of the directive.
+      directiveAnnot :: a -- ^ Annotation for the directive.
+    } -> Directive f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- A directive body has multiple directive elements.
+-- | Represents the body of a directive, which consists of multiple elements.
 data DirectiveBody :: SynTree where
-  DirectiveBody :: [DirectiveElement f a] -> a -> DirectiveBody f a
+  DirectiveBody ::
+    { directiveElements :: [DirectiveElement f a], -- ^ Elements of the directive.
+      directiveBodyAnnot :: a -- ^ Annotation for the directive body.
+    } -> DirectiveBody f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- Element in the directive.
+-- | Represents an element in a directive. Can be either a key or a key-value 
+--   pair.
 data DirectiveElement :: SynTree where
-  -- A directive can just be a key. Where the mere presence of the key has
-  -- semantic value.
+  -- | A simple directive element with a key. The mere presence of this key
+  --   holds semantic value.
   DirectiveElementKey ::
-    -- Which backend is this directive intended for?
-    Maybe (Identifier f a) ->
-    Identifier f a ->
-    a ->
-    DirectiveElement f a
-  -- A directive can be more complex too. It can have an optional backend
-  -- specificer, a key and a value.
+    { directiveBackend :: Maybe (Identifier f a), -- ^ Optional backend target.
+      directiveKey :: Identifier f a, -- ^ The key of the directive.
+      directiveKeyAnnot :: a -- ^ Annotation for the directive element.
+    } -> DirectiveElement f a
+  -- | A more complex directive element with a key-value pair, optionally
+  --   specifying a backend.
   DirectiveElementKeyValue ::
-    -- Which backend is this directive intendend for?
-    Maybe (Identifier f a) ->
-    -- The key for this directive.
-    Identifier f a ->
-    -- The value for this directive.
-    DirectiveExpression f a ->
-    a ->
-    DirectiveElement f a
+    { directiveBackend :: Maybe (Identifier f a), -- ^ Optional backend target.
+      directiveKey :: Identifier f a, -- ^ The key of the directive.
+      directiveValue :: DirectiveExpression f a, -- ^ The value of the directive.
+      directiveKeyValueAnnot :: a -- ^ Annotation for the key-value directive.
+    } -> DirectiveElement f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- Expressions which can be found in the directive.
+-- | Represents expressions that can be used within a directive, either a 
+--   string or a number.
 data DirectiveExpression f a where
-  DirectiveString :: Text -> a -> DirectiveExpression f a
-  DirectiveNumber :: Text -> a -> DirectiveExpression f a
+  DirectiveString :: 
+    { directiveStringValue :: Text, -- ^ String value of the directive.
+      directiveStringAnnot :: a -- ^ Annotation for the directive string.
+    } -> DirectiveExpression f a
+  DirectiveNumber :: 
+    { directiveNumberValue :: Text, -- ^ Number value of the directive.
+      directiveNumberAnnot :: a -- ^ Annotation for the directive number.
+    } -> DirectiveExpression f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- A type, which wraps another syntax tree, but tacks on an array of directives.
--- that apply to the subtree.
+-- | A type that wraps another syntax tree and applies a list of directives to 
+--   it.
 data Directed t stage f a where
-  Directed :: [Directive f a] -> t stage f a -> a -> Directed t stage f a
+  Directed ::
+    { directedDirectives :: [Directive f a], -- ^ List of directives.
+      directedSubtree :: t stage f a, -- ^ The wrapped syntax tree.
+      directedAnnot :: a -- ^ Annotation for the directed subtree.
+    } -> Directed t stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- Apply a function to the underlying subtree in a Directed type.
+-- | Apply a function to the underlying subtree in a 'Directed' type.
 mapDirected :: (t s f a -> t' s' f a) -> Directed t s f a -> Directed t' s' f a
 mapDirected fn (Directed dr tfa a) = Directed dr (fn tfa) a
 
--- Apply a monadic function to the underlying subtree in a Directed type.
+-- | Apply a monadic function to the underlying subtree in a 'Directed' type.
 mapDirectedM ::
-  (Monad m) => (t s f a -> m (t' s' f a)) -> Directed t s f a -> m (Directed t' s' f a)
+  (Monad m) => 
+  (t s f a -> m (t' s' f a)) -> 
+  Directed t s f a -> 
+  m (Directed t' s' f a)
 mapDirectedM fn (Directed dr tfa a) = Directed dr <$> fn tfa <*> pure a
 
+-- | Convert an annotated syntax tree element to a 'Directed' type with 
+--   an empty directive list.
 asDirected :: (Annotated (t s)) => t s f a -> Directed t s f a
 asDirected tfa = Directed [] tfa (annot tfa)
 
+-- | Extract the underlying subtree from a 'Directed' type, discarding any 
+--   directives.
 undirected :: Directed t s f a -> t s f a
 undirected (Directed _ tfa _) = tfa
 
--- Root of the parse tree. Just contains a list of declarations.
+-- | The root of the parse tree, containing a list of top-level declarations.
 data FiddleUnit (stage :: Stage) (f :: Type -> Type) a where
-  FiddleUnit :: [Directed FiddleDecl stage f a] -> a -> FiddleUnit stage f a
+  FiddleUnit :: 
+    { fiddleDecls :: [Directed FiddleDecl stage f a], -- ^ List of declarations.
+      fiddleUnitAnnot :: a -- ^ Annotation for the 'FiddleUnit'.
+    } -> FiddleUnit stage f a
   deriving (Generic, Annotated, Typeable)
 
 deriving instance (Alter (ImportType stage)) => Alter (FiddleUnit stage)
 
--- Just an identifier.
-data Identifier f a = Identifier !Text a
+-- | Represents an identifier with an associated annotation.
+data Identifier f a = Identifier 
+  { identifierName :: !Text, -- ^ The name of the identifier.
+    identifierAnnot :: a -- ^ Annotation for the identifier.
+  }
   deriving (Generic, Annotated, Alter, Typeable)
 
--- Expression.
+-- | Expressions used within Fiddle, including literals and variables.
 data Expression (s :: Stage) :: SynTree where
-  -- Just a string. Parsing the number comes in stage2.
-  LitNum :: NumberType stage -> a -> Expression stage f a
-  Var :: Identifier f a -> a -> Expression stage f a
+  -- | A numeric literal, whose value is dependent on the compilation stage.
+  LitNum :: 
+    { litNumValue :: NumberType stage, -- ^ The numeric value.
+      litNumAnnot :: a -- ^ Annotation for the literal.
+    } -> Expression stage f a
+  -- | A variable reference.
+  Var :: 
+    { varIdentifier :: Identifier f a, -- ^ The identifier of the variable.
+      varAnnot :: a -- ^ Annotation for the variable.
+    } -> Expression stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents an import statement in the Fiddle language.
 data ImportStatement f a where
-  ImportStatement :: Text -> Maybe (ImportList f a) -> a -> ImportStatement f a
+  ImportStatement :: 
+    { importPath :: Text, -- ^ The path to import.
+      importList :: Maybe (ImportList f a), -- ^ Optional list of imported items.
+      importStatementAnnot :: a -- ^ Annotation for the import statement.
+    } -> ImportStatement f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | A list of imported identifiers.
 data ImportList f a where
-  ImportList :: [Identifier f a] -> a -> ImportList f a
+  ImportList :: 
+    { importIdentifiers :: [Identifier f a], -- ^ The list of identifiers.
+      importListAnnot :: a -- ^ Annotation for the import list.
+    } -> ImportList f a
   deriving (Generic, Annotated, Alter, Typeable)
 
--- Top-level declarations.
+-- | Represents top-level declarations in Fiddle.
 data FiddleDecl :: StagedSynTree where
-  {-
-   - An option is a key/value pair.
-   - option <ident> <ident>;
-   -}
+  -- | An option declaration in the form 'option <ident> <ident>'.
   OptionDecl ::
-    Identifier f a ->
-    Identifier f a ->
-    a ->
-    FiddleDecl stage f a
+    { optionKey :: Identifier f a, -- ^ The key of the option.
+      optionValue :: Identifier f a, -- ^ The value of the option.
+      optionAnnot :: a -- ^ Annotation for the option declaration.
+    } -> FiddleDecl stage f a
+  -- | An import declaration.
   ImportDecl ::
-    ImportType stage f a ->
-    a ->
-    FiddleDecl stage f a
+    { importType :: ImportType stage f a, -- ^ The imported type.
+      importDeclAnnot :: a -- ^ Annotation for the import declaration.
+    } -> FiddleDecl stage f a
+  -- | A using declaration.
   UsingDecl ::
-    Name f a -> a -> FiddleDecl stage f a
-  {- Package Statement. Package Name, Package body -}
+    { usingName :: Name f a, -- ^ The name being used.
+      usingAnnot :: a -- ^ Annotation for the using declaration.
+    } -> FiddleDecl stage f a
+  -- | A package declaration.
   PackageDecl ::
-    Name f a ->
-    f (PackageBody stage f a) ->
-    a ->
-    FiddleDecl stage f a
-  {- location <identifier> = <expr>. -}
+    { packageName :: Name f a, -- ^ The package name.
+      packageBody :: f (PackageBody stage f a), -- ^ The body of the package.
+      packageAnnot :: a -- ^ Annotation for the package declaration.
+    } -> FiddleDecl stage f a
+  -- | A location declaration in the form 'location <identifier> = <expr>'.
   LocationDecl ::
-    Identifier f a ->
-    Expression stage f a ->
-    a ->
-    FiddleDecl stage f a
-  {- bits <identifier> : <type> -}
+    { locationIdent :: Identifier f a, -- ^ The location identifier.
+      locationExpr :: Expression stage f a, -- ^ The associated expression.
+      locationAnnot :: a -- ^ Annotation for the location declaration.
+    } -> FiddleDecl stage f a
+  -- | A bits declaration in the form 'bits <identifier> : <type>'.
   BitsDecl ::
-    Identifier f a ->
-    BitType stage f a ->
-    a ->
-    FiddleDecl stage f a
-  {- objtype <identifier> : <type> -}
+    { bitsIdent :: Identifier f a, -- ^ The identifier of the bits.
+      bitsType :: BitType stage f a, -- ^ The type of the bits.
+      bitsAnnot :: a -- ^ Annotation for the bits declaration.
+    } -> FiddleDecl stage f a
+  -- | An object type declaration.
   ObjTypeDecl ::
-    Identifier f a ->
-    f (ObjTypeBody stage f a) ->
-    a ->
-    FiddleDecl stage f a
-  {- object <ident> at <expr> : <type> -}
+    { objTypeIdent :: Identifier f a, -- ^ The identifier of the object type.
+      objTypeBody :: f (ObjTypeBody stage f a), -- ^ The body of the object type.
+      objTypeAnnot :: a -- ^ Annotation for the object type declaration.
+    } -> FiddleDecl stage f a
+  -- | An object declaration in the form 'object <ident> at <expr> : <type>'.
   ObjectDecl ::
-    Identifier f a ->
-    Expression stage f a ->
-    ObjType stage f a ->
-    a ->
-    FiddleDecl stage f a
+    { objectIdent :: Identifier f a, -- ^ The identifier of the object.
+      objectLocation :: Expression stage f a, -- ^ The location expression.
+      objectType :: ObjType stage f a, -- ^ The type of the object.
+      objectAnnot :: a -- ^ Annotation for the object declaration.
+    } -> FiddleDecl stage f a
   deriving (Generic, Annotated, Typeable)
 
 deriving instance (Alter (ImportType stage)) => Alter (FiddleDecl stage)
 
+-- | Represents the body of an object type, containing a body type (struct or
+--   union), a list of object declarations, and an annotation.
 data ObjTypeBody (stage :: Stage) (f :: Type -> Type) a where
   ObjTypeBody ::
-    BodyType f a ->
-    [Directed ObjTypeDecl stage f a] ->
-    a ->
-    ObjTypeBody stage f a
+    { objBodyType :: BodyType f a, -- ^ The body type (struct or union).
+      objBodyDecls :: [Directed ObjTypeDecl stage f a], -- ^ Object declarations.
+      objBodyAnnot :: a -- ^ Annotation for the object type body.
+    } -> ObjTypeBody stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents an object type, which can be anonymous, an array, or a
+--   reference to another type.
 data ObjType stage f a where
-  -- { <body> }
-  -- Anonymous types are only allowed in stage1. Stage2 should have them be
-  -- de-anonymized.
+  -- | An anonymous object type, allowed only in Stage1. 
   AnonymousObjType ::
-    (Witness (stage == Stage1)) ->
-    f (ObjTypeBody stage f a) ->
-    a ->
-    ObjType stage f a
-  -- <type>[<expr>]
-  ArrayObjType :: ObjType stage f a -> Expression stage f a -> a -> ObjType stage f a
-  -- <identifier>
-  ReferencedObjType :: Name f a -> a -> ObjType stage f a
+    { anonWitness :: Witness (stage == Stage1), -- ^ Witness for stage constraint.
+      anonBody :: f (ObjTypeBody stage f a), -- ^ The body of the anonymous type.
+      anonAnnot :: a -- ^ Annotation for the anonymous type.
+    } -> ObjType stage f a
+  -- | An array of object types.
+  ArrayObjType ::
+    { arrayObjType :: ObjType stage f a, -- ^ The type of the array elements.
+      arraySize :: Expression stage f a, -- ^ The size of the array.
+      arrayAnnot :: a -- ^ Annotation for the array type.
+    } -> ObjType stage f a
+  -- | A reference to an existing type by name.
+  ReferencedObjType ::
+    { refName :: Name f a, -- ^ The name of the referenced type.
+      refAnnot :: a -- ^ Annotation for the referenced type.
+    } -> ObjType stage f a
   deriving (Typeable, Generic, Alter, Annotated, Typeable)
 
+-- | Represents a declaration inside an object type, such as a register, an
+--   assertion, or a substructure.
 data ObjTypeDecl stage f a where
-  {- assert_pos(<expr>) -}
+  -- | An assertion statement for a specific position.
   AssertPosStatement ::
-    Witness (StageLessThan stage 3) ->
-    Expression stage f a ->
-    a ->
-    ObjTypeDecl stage f a
-  {- reg <ident>(<expr>) : <regtype> -}
+    { assertWitness :: Witness (StageLessThan stage 3), -- ^ Witness for stage constraint.
+      assertExpr :: Expression stage f a, -- ^ The expression for the assertion.
+      assertAnnot :: a -- ^ Annotation for the assertion.
+    } -> ObjTypeDecl stage f a
+  -- | A register declaration.
   RegisterDecl ::
-    Maybe (Modifier f a) ->
-    Maybe (Identifier f a) ->
-    Expression stage f a ->
-    Maybe (RegisterBody stage f a) ->
-    a ->
-    ObjTypeDecl stage f a
-  {- reserved(n); -}
-  ReservedDecl :: Expression stage f a -> a -> ObjTypeDecl stage f a
-  {- <struct|union> { subfields } <name>; -}
+    { regModifier :: Maybe (Modifier f a), -- ^ Optional register modifier.
+      regIdent :: Maybe (Identifier f a), -- ^ Optional register identifier.
+      regSize :: Expression stage f a, -- ^ Size of the register.
+      regBody :: Maybe (RegisterBody stage f a), -- ^ Optional register body.
+      regAnnot :: a -- ^ Annotation for the register declaration.
+    } -> ObjTypeDecl stage f a
+  -- | A reserved declaration for padding or alignment.
+  ReservedDecl ::
+    { reservedExpr :: Expression stage f a, -- ^ The expression for reserved space.
+      reservedAnnot :: a -- ^ Annotation for the reserved declaration.
+    } -> ObjTypeDecl stage f a
+  -- | A declaration for a substructure (struct or union).
   TypeSubStructure ::
-    f (ObjTypeBody stage f a) ->
-    Maybe (Identifier f a) ->
-    a ->
-    ObjTypeDecl stage f a
+    { subStructureBody :: f (ObjTypeBody stage f a), -- ^ The body of the substructure.
+      subStructureName :: Maybe (Identifier f a), -- ^ Optional name for the substructure.
+      subStructureAnnot :: a -- ^ Annotation for the substructure.
+    } -> ObjTypeDecl stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents a modifier for registers (e.g., read-only, read-write).
 data Modifier f a where
-  ModifierKeyword :: ModifierKeyword -> a -> Modifier f a
+  ModifierKeyword ::
+    { modifierKey :: ModifierKeyword, -- ^ The keyword for the modifier.
+      modifierAnnot :: a -- ^ Annotation for the modifier.
+    } -> Modifier f a
   deriving (Generic, Annotated, Alter, Typeable)
 
-data ModifierKeyword = Rw | Ro | Wo deriving (Eq, Ord, Show, Read, Typeable)
+-- | Enumerates the different types of register modifiers.
+data ModifierKeyword = Rw | Ro | Wo
+  deriving (Eq, Ord, Show, Read, Typeable)
 
+-- | Represents a deferred register body, consisting of a list of bit
+--   declarations.
 data DeferredRegisterBody stage f a where
   DeferredRegisterBody ::
-    [Directed RegisterBitsDecl stage f a] ->
-    a ->
-    DeferredRegisterBody stage f a
+    { deferredBits :: [Directed RegisterBitsDecl stage f a], -- ^ Bit declarations.
+      deferredAnnot :: a -- ^ Annotation for the deferred register body.
+    } -> DeferredRegisterBody stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents the body type (struct or union) in an object.
 data BodyType (f :: Type -> Type) a where
-  Union :: a -> BodyType f a
-  Struct :: a -> BodyType f a
+  Union ::
+    { unionAnnot :: a -- ^ Annotation for the union.
+    } -> BodyType f a
+  Struct ::
+    { structAnnot :: a -- ^ Annotation for the struct.
+    } -> BodyType f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents a register body with a body type and deferred bit declarations.
 data RegisterBody stage f a where
   RegisterBody ::
-    BodyType f a ->
-    f (DeferredRegisterBody stage f a) ->
-    a ->
-    RegisterBody stage f a
+    { regBodyType :: BodyType f a, -- ^ The body type of the register.
+      regDeferredBody :: f (DeferredRegisterBody stage f a), -- ^ Deferred body.
+      regBodyAnnot :: a -- ^ Annotation for the register body.
+    } -> RegisterBody stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents declarations within a register, such as defined bits,
+--   reserved bits, or substructures.
 data RegisterBitsDecl stage f a where
-  -- reserved(<expr>)
-  ReservedBits :: Expression stage f a -> a -> RegisterBitsDecl stage f a
-  -- <modifer> <ident> : <type>
+  -- | Declaration for reserved bits.
+  ReservedBits ::
+    { reservedBitsExpr :: Expression stage f a, -- ^ Expression for reserved bits.
+      reservedBitsAnnot :: a -- ^ Annotation for the reserved bits.
+    } -> RegisterBitsDecl stage f a
+  -- | Declaration for defined bits in a register.
   DefinedBits ::
-    Maybe (Modifier f a) ->
-    Identifier f a ->
-    RegisterBitsTypeRef stage f a ->
-    a ->
-    RegisterBitsDecl stage f a
+    { definedBitsModifier :: Maybe (Modifier f a), -- ^ Optional modifier for the bits.
+      definedBitsIdent :: Identifier f a, -- ^ Identifier for the bits.
+      definedBitsTypeRef :: RegisterBitsTypeRef stage f a, -- ^ Type reference for the bits.
+      definedBitsAnnot :: a -- ^ Annotation for the defined bits.
+    } -> RegisterBitsDecl stage f a
+  -- | Substructure within a register.
   BitsSubStructure ::
-    RegisterBody stage f a ->
-    Maybe (Identifier f a) ->
-    a ->
-    RegisterBitsDecl stage f a
+    { bitsSubRegBody :: RegisterBody stage f a, -- ^ The body of the substructure.
+      bitsSubName :: Maybe (Identifier f a), -- ^ Optional name for the substructure.
+      bitsSubAnnot :: a -- ^ Annotation for the substructure.
+    } -> RegisterBitsDecl stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents different ways to refer to register bits, either as an array,
+--   a reference to a type, an anonymous type, or just bits.
 data RegisterBitsTypeRef stage f a where
-  -- <type>[<expr>]
+  -- | An array of bits with a specified size.
   RegisterBitsArray ::
-    RegisterBitsTypeRef stage f a ->
-    Expression stage f a ->
-    a ->
-    RegisterBitsTypeRef stage f a
-  {- Reference to a type. -}
-  RegisterBitsReference :: Name f a -> a -> RegisterBitsTypeRef stage f a
-  {- enum(<expr>) { <body> }
-     Anonymous types are only allowed in stage1.
-     Stage2 should de-anonymize these type. -}
+    { bitsArrayTypeRef :: RegisterBitsTypeRef stage f a, -- ^ Reference to the array type.
+      bitsArraySize :: Expression stage f a, -- ^ Size of the array.
+      bitsArrayAnnot :: a -- ^ Annotation for the array.
+    } -> RegisterBitsTypeRef stage f a
+  -- | A reference to another type by name.
+  RegisterBitsReference ::
+    { bitsRefName :: Name f a, -- ^ The name of the referenced type.
+      bitsRefAnnot :: a -- ^ Annotation for the reference.
+    } -> RegisterBitsTypeRef stage f a
+  -- | An anonymous type for register bits, used in Stage1.
   RegisterBitsAnonymousType ::
-    Witness (stage == Stage1) ->
-    AnonymousBitsType stage f a ->
-    a ->
-    RegisterBitsTypeRef stage f a
-  {- (<expr>)
-   -
-   - The expression is just bits ... i.e. an integer.
-   -}
+    { anonBitsWitness :: Witness (stage == Stage1), -- ^ Witness for stage constraint.
+      anonBitsType :: AnonymousBitsType stage f a, -- ^ The anonymous type.
+      anonBitsAnnot :: a -- ^ Annotation for the anonymous type.
+    } -> RegisterBitsTypeRef stage f a
+  -- | A direct specification of bits as an expression.
   RegisterBitsJustBits ::
-    Expression stage f a ->
-    a ->
-    RegisterBitsTypeRef stage f a
+    { justBitsExpr :: Expression stage f a, -- ^ Expression for the bits.
+      justBitsAnnot :: a -- ^ Annotation for the bits.
+    } -> RegisterBitsTypeRef stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents an anonymous bit type, such as an enum, used in Stage1.
 data AnonymousBitsType stage f a where
-  -- enum(<expr>) { <body> }
   AnonymousEnumBody ::
-    Expression stage f a ->
-    f (EnumBody stage f a) ->
-    a ->
-    AnonymousBitsType stage f a
+    { anonEnumExpr :: Expression stage f a, -- ^ Expression defining the enum size.
+      anonEnumBody :: f (EnumBody stage f a), -- ^ The body of the enum.
+      anonEnumAnnot :: a -- ^ Annotation for the anonymous enum.
+    } -> AnonymousBitsType stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents a bit type, either an enumeration or raw bits.
 data BitType (stage :: Stage) (f :: Type -> Type) a where
-  -- enum(<expr>) { <body> }
+  -- | An enumeration type for bits.
   EnumBitType ::
-    Expression stage f a ->
-    f (EnumBody stage f a) ->
-    a ->
-    BitType stage f a
-  -- (<expr>)
-  RawBits :: Expression stage f a -> a -> BitType stage f a
+    { enumBitExpr :: Expression stage f a, -- ^ Expression defining the enum size.
+      enumBitBody :: f (EnumBody stage f a), -- ^ The body of the enum.
+      enumBitAnnot :: a -- ^ Annotation for the enumeration.
+    } -> BitType stage f a
+  -- | A raw bit type.
+  RawBits ::
+    { rawBitsExpr :: Expression stage f a, -- ^ Expression defining the bits.
+      rawBitsAnnot :: a -- ^ Annotation for the raw bits.
+    } -> BitType stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents the body of an enumeration.
 data EnumBody (stage :: Stage) (f :: Type -> Type) a where
-  -- <decl>,
-  EnumBody :: [Directed EnumConstantDecl stage f a] -> a -> EnumBody stage f a
+  EnumBody ::
+    { enumConsts :: [Directed EnumConstantDecl stage f a], -- ^ Enum constant declarations.
+      enumBodyAnnot :: a -- ^ Annotation for the enum body.
+    } -> EnumBody stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents a declaration for an enumeration constant.
 data EnumConstantDecl stage f a where
-  -- <ident> = <expr>
+  -- | A named constant in the enum.
   EnumConstantDecl ::
-    Identifier f a ->
-    Expression stage f a ->
-    a ->
-    EnumConstantDecl stage f a
-  -- reserved = <expr>
+    { enumConstIdent :: Identifier f a, -- ^ Identifier for the constant.
+      enumConstExpr :: Expression stage f a, -- ^ Expression defining the constant.
+      enumConstAnnot :: a -- ^ Annotation for the constant.
+    } -> EnumConstantDecl stage f a
+  -- | A reserved value in the enum.
   EnumConstantReserved ::
-    Expression stage f a ->
-    a ->
-    EnumConstantDecl stage f a
+    { enumReservedExpr :: Expression stage f a, -- ^ Expression for the reserved value.
+      enumReservedAnnot :: a -- ^ Annotation for the reserved value.
+    } -> EnumConstantDecl stage f a
   deriving (Generic, Annotated, Alter, Typeable)
 
+-- | Represents the body of a package, containing a list of declarations.
 data PackageBody (stage :: Stage) (f :: Type -> Type) a where
-  {- The body of a package -}
-  PackageBody :: [Directed FiddleDecl stage f a] -> a -> PackageBody stage f a
+  PackageBody ::
+    { packageBodyDecls :: [Directed FiddleDecl stage f a], -- ^ Declarations in the package.
+      packageBodyAnnot :: a -- ^ Annotation for the package body.
+    } -> PackageBody stage f a
   deriving (Generic, Annotated, Typeable)
 
 deriving instance (Alter (ImportType stage)) => Alter (PackageBody stage)
 
+
 squeeze :: (Alter t, Traversable f, Monad f) => t f a -> f (t Identity a)
 squeeze = alter (fmap Identity) return
-
--- Expression involves NumberType, so we add the constraint:
-deriving instance (Functor f, NumberType s ~ NumberType s') => EasySwitchStage Expression f s s'
-
--- FiddleDecl includes both NumberType and ImportType, so we need both constraints:
-deriving instance
-  ( Functor f,
-    StageLessThan s 3 ~ StageLessThan s' 3,
-    (s == Stage1) ~ (s' == Stage1),
-    NumberType s ~ NumberType s',
-    ImportType s ~ ImportType s'
-  ) =>
-  EasySwitchStage FiddleDecl f s s'
-
--- ObjType includes NumberType, so we add the constraint:
-deriving instance
-  ( Functor f,
-    StageLessThan s 3 ~ StageLessThan s' 3,
-    (s == Stage1) ~ (s' == Stage1),
-    NumberType s ~ NumberType s'
-  ) =>
-  EasySwitchStage ObjType f s s'
-
--- ObjTypeBody doesn't have any special type families, so no additional constraints:
-deriving instance
-  ( Functor f,
-    StageLessThan s 3 ~ StageLessThan s' 3,
-    (s == Stage1) ~ (s' == Stage1),
-    NumberType s ~ NumberType s'
-  ) =>
-  EasySwitchStage ObjTypeBody f s s'
-
--- ObjTypeDecl doesn't have special type families, so no additional constraints:
-deriving instance
-  ( Functor f,
-    StageLessThan s 3 ~ StageLessThan s' 3,
-    (s == Stage1) ~ (s' == Stage1),
-    NumberType s ~ NumberType s'
-  ) =>
-  EasySwitchStage ObjTypeDecl f s s'
-
--- DeferredRegisterBody doesn't have special type families:
-deriving instance
-  ( Functor f,
-    (s == Stage1) ~ (s' == Stage1),
-    NumberType s ~ NumberType s'
-  ) =>
-  EasySwitchStage DeferredRegisterBody f s s'
-
--- RegisterBody depends on NumberType, so we add that constraint:
-deriving instance
-  ( Functor f,
-    (s == Stage1) ~ (s' == Stage1),
-    NumberType s ~ NumberType s'
-  ) =>
-  EasySwitchStage RegisterBody f s s'
-
--- RegisterBitsDecl depends on NumberType:
-deriving instance
-  ( Functor f,
-    (s == Stage1) ~ (s' == Stage1),
-    NumberType s ~ NumberType s'
-  ) =>
-  EasySwitchStage RegisterBitsDecl f s s'
-
--- RegisterBitsTypeRef depends on NumberType:
-deriving instance
-  ( Functor f,
-    (s == Stage1) ~ (s' == Stage1),
-    NumberType s ~ NumberType s'
-  ) =>
-  EasySwitchStage RegisterBitsTypeRef f s s'
-
--- AnonymousBitsType depends on NumberType:
-deriving instance
-  (Functor f, NumberType s ~ NumberType s') =>
-  EasySwitchStage AnonymousBitsType f s s'
-
--- BitType depends on NumberType:
-deriving instance
-  (Functor f, NumberType s ~ NumberType s') =>
-  EasySwitchStage BitType f s s'
-
--- EnumBody doesn't depend on any type families:
-deriving instance
-  (Functor f, NumberType s ~ NumberType s') =>
-  EasySwitchStage EnumBody f s s'
-
--- EnumConstantDecl depends on NumberType:
-deriving instance
-  (Functor f, NumberType s ~ NumberType s') =>
-  EasySwitchStage EnumConstantDecl f s s'
-
--- PackageBody includes both NumberType and ImportType:
-deriving instance
-  ( Functor f,
-    StageLessThan s 3 ~ StageLessThan s' 3,
-    (s == Stage1) ~ (s' == Stage1),
-    NumberType s ~ NumberType s',
-    ImportType s ~ ImportType s'
-  ) =>
-  EasySwitchStage PackageBody f s s'
-
--- FiddleUnit includes NumberType and ImportType, so we need both constraints:
-deriving instance
-  ( Functor f,
-    StageLessThan s 3 ~ StageLessThan s' 3,
-    (s == Stage1) ~ (s' == Stage1),
-    NumberType s ~ NumberType s',
-    ImportType s ~ ImportType s'
-  ) =>
-  EasySwitchStage FiddleUnit f s s'
-
--- Directed depends on its underlying AST type:
-deriving instance
-  (EasySwitchStage t f s s') =>
-  EasySwitchStage (Directed t) f s s'
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
diff --git a/src/Language/Fiddle/GenericTree.hs b/src/Language/Fiddle/GenericTree.hs
index 59db6aa..031f6ab 100644
--- a/src/Language/Fiddle/GenericTree.hs
+++ b/src/Language/Fiddle/GenericTree.hs
@@ -40,17 +40,38 @@ type Context stage =
 data GenericSyntaxTree f a where
   {- GenericSyntaxtTree with a name and children. -}
   SyntaxTreeObject ::
-    forall a f tree.
-    (Typeable tree) =>
+    forall t f a.
+    (Typeable t, Typeable f, Typeable a) =>
     String ->
     [GenericSyntaxTree f a] ->
     a ->
-    tree ->
+    t f a ->
     GenericSyntaxTree f a
   SyntaxTreeList :: [GenericSyntaxTree f a] -> GenericSyntaxTree f a
   SyntaxTreeDeferred :: f (GenericSyntaxTree f a) -> GenericSyntaxTree f a
   SyntaxTreeValue :: String -> GenericSyntaxTree f a
 
+alterGenericSyntaxTree ::
+  (Functor f) =>
+  (GenericSyntaxTree f a -> Maybe (GenericSyntaxTree f a)) ->
+  GenericSyntaxTree f a ->
+  GenericSyntaxTree f a
+alterGenericSyntaxTree fn genericTree
+  | (Just newGenericTree) <- fn genericTree = newGenericTree
+  | otherwise =
+      case genericTree of
+        SyntaxTreeObject str members annot tree ->
+          SyntaxTreeObject
+            str
+            (map (alterGenericSyntaxTree fn) members)
+            annot
+            tree
+        SyntaxTreeList members ->
+          SyntaxTreeList $ map (alterGenericSyntaxTree fn) members
+        SyntaxTreeDeferred sub ->
+          SyntaxTreeDeferred $ fmap (alterGenericSyntaxTree fn) sub
+        v -> v
+
 instance ToJSON Comment where
   toJSON (NormalComment str) = object ["normal" .= str]
   toJSON (DocComment str) = object ["doc" .= str]
@@ -71,9 +92,11 @@ instance ToJSON SourcePos where
         "col" .= sourceColumn sourcePos
       ]
 
-instance (Foldable f, ToJSON a) => ToJSON (GenericSyntaxTree f a) where
+instance (Foldable f) => ToJSON (GenericSyntaxTree f (Maybe Value)) where
   toJSON = \case
-    (SyntaxTreeObject typ membs a _) ->
+    (SyntaxTreeObject typ membs Nothing _) ->
+      object ["_con" .= typ, "_members" .= membs]
+    (SyntaxTreeObject typ membs (Just a) _) ->
       object ["_con" .= typ, "_members" .= membs, "_annot" .= a]
     (SyntaxTreeList l) ->
       Array $ Data.Vector.fromList $ map toJSON l
@@ -82,8 +105,10 @@ instance (Foldable f, ToJSON a) => ToJSON (GenericSyntaxTree f a) where
     (SyntaxTreeValue s) -> String (Data.Text.pack s)
 
   toEncoding = \case
-    (SyntaxTreeObject typ membs a t) ->
+    (SyntaxTreeObject typ membs (Just a) t) ->
       pairs $ "_type" .= show (typeOf t) <> "_con" .= typ <> "_members" .= membs <> "_annot" .= a
+    (SyntaxTreeObject typ membs Nothing t) ->
+      pairs $ "_type" .= show (typeOf t) <> "_con" .= typ <> "_members" .= membs
     (SyntaxTreeList l) ->
       foldable $ map toJSON l
     (SyntaxTreeDeferred fdef) ->
@@ -118,7 +143,7 @@ class ToGenericSyntaxTree (t :: SynTree) where
   toGenericSyntaxTree t = gToGenericSyntaxTree t (from t)
 
 class GToGenericSyntaxTree r f a where
-  gToGenericSyntaxTree :: (Typeable t) => t -> r x -> GenericSyntaxTree f a
+  gToGenericSyntaxTree :: (Typeable t, Typeable f, Typeable a) => t f a -> r x -> GenericSyntaxTree f a
 
 class GToMemberList r f a where
   gToMemberList :: Int -> r x -> [GenericSyntaxTree f a]
diff --git a/src/Language/Fiddle/Internal/Scopes.hs b/src/Language/Fiddle/Internal/Scopes.hs
new file mode 100644
index 0000000..280945d
--- /dev/null
+++ b/src/Language/Fiddle/Internal/Scopes.hs
@@ -0,0 +1,101 @@
+module Language.Fiddle.Internal.Scopes where
+
+import Control.Monad (forM)
+import Data.List (inits, intercalate)
+import Data.List.NonEmpty (NonEmpty (..), prependList)
+import Data.Map (Map)
+import qualified Data.Map as Map
+import Data.Maybe (fromMaybe)
+
+-- | 'Scope' represents a hierarchical structure for storing key-value pairs.
+--   It can contain nested sub-scopes, which are stored in 'subScopes',
+--   and the values for a specific scope are stored in 'scopeValues'.
+data Scope k v = Scope
+  { subScopes :: Map k (Scope k v), -- Nested sub-scopes
+    scopeValues :: Map k v -- Values stored in the current scope
+  }
+
+-- | 'ScopePath' keeps track of the current scope path as a list of keys,
+--   and also includes any additional paths (like imported modules or
+--   using namespaces) that might be referenced for lookup.
+data ScopePath k = ScopePath
+  { currentScope :: [k], -- Current path within the scope hierarchy
+    usingPaths :: [[k]] -- Additional paths for resolving symbols
+  }
+
+-- | The 'Semigroup' instance for 'Scope' allows combining two scopes,
+--   where sub-scopes and values are merged together.
+instance (Ord k) => Semigroup (Scope k t) where
+  (Scope ss1 sv1) <> (Scope ss2 sv2) =
+    Scope (Map.unionWith (<>) ss1 ss2) (Map.union sv1 sv2)
+
+-- | The 'Monoid' instance for 'Scope' provides an empty scope with
+--   no sub-scopes or values.
+instance (Ord k) => Monoid (Scope k t) where
+  mempty = Scope mempty mempty
+
+-- | The 'Semigroup' instance for 'ScopePath' allows combining paths,
+--   appending the current scope path and using paths.
+instance Semigroup (ScopePath k) where
+  (ScopePath a1 b1) <> (ScopePath a2 b2) = ScopePath (a1 <> a2) (b1 <> b2)
+
+-- | The 'Monoid' instance for 'ScopePath' provides an empty path.
+instance Monoid (ScopePath k) where
+  mempty = ScopePath mempty mempty
+
+-- | 'insertScope' inserts a value 'v' into the scope at the specified
+--   key path ('NonEmpty k'). If the key path does not exist, it is created.
+insertScope :: (Ord k) => NonEmpty k -> t -> Scope k t -> Scope k 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 mempty -> mp) -> Just (insertScope (a :| as) v mp)
+        )
+        s
+        ss
+    )
+    sv
+
+-- | 'lookupScope' performs a lookup of a value in the scope using a key path
+--   ('NonEmpty k'). It traverses through sub-scopes as defined by the path.
+lookupScope :: (Ord k) => NonEmpty k -> Scope k 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' searches for a key in the scope by trying all possible
+--   paths, including the current scope path and any additional 'using' paths.
+--   It returns a list of all valid matches, each paired with the corresponding
+--   full key path that resolved to a value. This is useful in cases where symbol
+--   resolution might be ambiguous or multiple valid resolutions exist.
+--
+--   The result is a list of tuples, where the first element of each tuple is the
+--   fully-resolved key path (as a 'NonEmpty' list) and the second element is the
+--   resolved value.
+--
+--   This function handles multiple levels of scope resolution, such as those
+--   introduced by 'using' directives, and concatenates the results from each
+--   possible path.
+--
+--   Parameters:
+--   * 'ScopePath k' - The current scope path and any additional 'using' paths.
+--   * 'NonEmpty k'  - The key path to resolve within the scope.
+--   * 'Scope k t'   - The scope containing sub-scopes and values.
+--
+--   Returns: A list of tuples where each tuple contains a fully-resolved key
+--   path (as a 'NonEmpty k') and the corresponding value ('t').
+lookupScopeWithPath ::
+  (Ord k) =>
+  ScopePath k ->
+  NonEmpty k ->
+  Scope k t ->
+  [(NonEmpty k, t)]
+lookupScopeWithPath (ScopePath current others) key scope =
+  let allPaths = reverse (inits current) ++ others
+   in flip concatMap allPaths $ \prefix -> do
+        case lookupScope (prependList prefix key) scope of
+          Just s -> [(prependList prefix key, s)]
+          Nothing -> []
diff --git a/src/Language/Fiddle/Parser.hs b/src/Language/Fiddle/Parser.hs
index 85ae65e..980925f 100644
--- a/src/Language/Fiddle/Parser.hs
+++ b/src/Language/Fiddle/Parser.hs
@@ -33,8 +33,8 @@ type Pa (a :: Stage -> (Type -> Type) -> Type -> Type) = P (a 'Stage1 F (Comment
 
 type PaS (a :: (Type -> Type) -> Type -> Type) = P (a F (Commented SourceSpan))
 
-comment :: P Comment
-comment =
+commentP :: P Comment
+commentP =
   token $ \case
     (TokComment c) -> Just (NormalComment c)
     (TokDocComment c) -> Just (DocComment c)
@@ -52,46 +52,46 @@ isComment (Token t _) =
 stripTrailingComments :: [Token s] -> [Token s]
 stripTrailingComments = reverse . dropWhile isComment . reverse
 
-directed :: Pa t -> PaS (Directed t 'Stage1)
-directed subparser = withMeta $ do
-  Directed <$> many directive <*> subparser
+directedP :: Pa t -> PaS (Directed t 'Stage1)
+directedP subparser = withMeta $ do
+  Directed <$> many directiveP <*> subparser
 
-directive :: PaS Directive
-directive =
+directiveP :: PaS Directive
+directiveP =
   withMeta $
-    Directive <$> defer directiveBodyTokens directiveBody
+    Directive <$> defer directiveBodyTokens directiveBodyP 
 
-directiveBody :: PaS DirectiveBody
-directiveBody = withMeta $ do
-  DirectiveBody <$> many (directiveElement <* (void (tok TokComma) <|> eof))
+directiveBodyP :: PaS DirectiveBody
+directiveBodyP = withMeta $ do
+  DirectiveBody <$> many (directiveElementP <* (void (tok TokComma) <|> eof))
 
-directiveElement :: PaS DirectiveElement
-directiveElement = withMeta $ do
-  identifier1 <- nextText
+directiveElementP :: PaS DirectiveElement
+directiveElementP = withMeta $ do
+  identifier1 <- nextTextP
   choice
     [ do
         tok TokColon
         let backend = identifier1
-        key <- nextText
+        key <- nextTextP
         choice
           [ do
               tok TokEq
-              DirectiveElementKeyValue (Just backend) key <$> directiveExpression,
+              DirectiveElementKeyValue (Just backend) key <$> directiveExpressionP,
             do
               return (DirectiveElementKey (Just backend) key)
           ],
       do
         tok TokEq
         let key = identifier1
-        DirectiveElementKeyValue Nothing key <$> directiveExpression,
+        DirectiveElementKeyValue Nothing key <$> directiveExpressionP,
       return $ DirectiveElementKey Nothing identifier1
     ]
 
-nextText :: PaS Identifier
-nextText = withMeta $ Identifier <$> token textOf
+nextTextP :: PaS Identifier
+nextTextP = withMeta $ Identifier <$> token textOf
 
-directiveExpression :: PaS DirectiveExpression
-directiveExpression = withMeta $ do
+directiveExpressionP :: PaS DirectiveExpression
+directiveExpressionP = withMeta $ do
   choice
     [ do
         token $ \case
@@ -103,80 +103,64 @@ directiveExpression = withMeta $ do
 fiddleUnit :: Pa FiddleUnit
 fiddleUnit = do
   withMeta
-    ( FiddleUnit <$> many1 (directed fiddleDecl <* tok TokSemi)
+    ( FiddleUnit <$> many1 (directedP fiddleDeclP <* tok TokSemi)
     )
-    <* many comment
+    <* many commentP
 
-stringToken :: P Text
-stringToken =
+stringTokenP :: P Text
+stringTokenP =
   token
     ( \case
         (TokString str) -> Just str
         _ -> Nothing
     )
 
-importList :: PaS ImportList
-importList = withMeta $ do
+importListP :: PaS ImportList
+importListP = withMeta $ do
   tok TokLParen
   ImportList
     <$> many (ident <* (tok TokComma <|> lookAhead (tok TokRParen)))
     <* tok TokRParen
 
-importStatement :: PaS ImportStatement
-importStatement =
+importStatementP :: PaS ImportStatement
+importStatementP =
   withMeta $
-    ImportStatement <$> stringToken <*> optionMaybe importList
+    ImportStatement <$> stringTokenP <*> optionMaybe importListP
 
-fiddleDecl :: Pa FiddleDecl
-fiddleDecl = do
+fiddleDeclP :: Pa FiddleDecl
+fiddleDeclP = do
   withMeta $ do
     t <- tokenType <$> anyToken
     case t of
-      KWOption -> OptionDecl <$> nextText <*> nextText
+      KWOption -> OptionDecl <$> nextTextP <*> nextTextP
       KWPackage ->
         PackageDecl
           <$> name
-          <*> defer body packageBody
+          <*> defer body packageBodyP
       KWUsing -> UsingDecl <$> name
-      KWLocation -> LocationDecl <$> ident <*> (tok TokEq >> expression)
-      KWBits -> BitsDecl <$> ident <*> (tok TokColon >> bitType)
-      KWImport -> ImportDecl <$> importStatement
+      KWLocation -> LocationDecl <$> ident <*> (tok TokEq >> expressionP)
+      KWBits -> BitsDecl <$> ident <*> (tok TokColon >> bitTypeP)
+      KWImport -> ImportDecl <$> importStatementP
       KWType ->
         ObjTypeDecl
           <$> ident
           <*> ( do
                   tok TokColon
-                  bt <- bodyType
-                  defer body (objTypeBody bt)
+                  bt <- bodyTypeP
+                  defer body (objTypeBodyP bt)
               )
       KWInstance ->
         ObjectDecl
           <$> ident
-          <*> (tok KWAt *> expression)
-          <*> (tok TokColon *> objType)
+          <*> (tok KWAt *> expressionP)
+          <*> (tok TokColon *> objTypeP)
       _ ->
         fail $
           printf "Unexpected token %s. Expected top-level declaration." (show t)
 
--- choice
---   [ tok KWOption >> OptionDecl <$> ident <*> ident,
---     tok KWPackage >> PackageDecl
---       <$> ident
---       <*> defer body packageBody,
---     tok KWLocation >> LocationDecl <$> ident <*> (tok TokEq >> expression),
---     tok KWBits >> BitsDecl <$> ident <*> (tok TokColon >> bitType),
---     tok KWObjtype
---       >> ObjTypeDecl <$> ident <*> (tok TokColon >> defer body objTypeBody),
---     tok KWObject
---       >> ObjectDecl
---         <$> ident
---         <*> (tok KWAt *> expression)
---         <*> (tok TokColon *> objType)
---   ]
-
-objType :: Pa ObjType
-objType = do
-  base <- withMeta baseObj
+objTypeP :: Pa ObjType
+objTypeP = do
+  base <- withMeta baseObjP
   recur' <- recur
   return $ recur' base
   where
@@ -184,54 +168,54 @@ objType = do
     recur =
       ( do
           withMeta $ do
-            expr <- tok TokLBracket *> expression <* tok TokRBracket
+            expr <- tok TokLBracket *> expressionP <* tok TokRBracket
             recur' <- recur
             return (\met base -> recur' (ArrayObjType base expr met))
       )
         <|> return id
 
-    baseObj :: P (A -> ObjType Stage1 F A)
-    baseObj =
+    baseObjP :: P (A -> ObjType Stage1 F A)
+    baseObjP =
       (ReferencedObjType <$> name)
         <|> ( do
-                t <- bodyType
-                AnonymousObjType (Witness ()) <$> defer body (objTypeBody t)
+                t <- bodyTypeP
+                AnonymousObjType (Witness ()) <$> defer body (objTypeBodyP t)
             )
 
-exprInParen :: Pa Expression
-exprInParen = tok TokLParen *> expression <* tok TokRParen
+exprInParenP :: Pa Expression
+exprInParenP = tok TokLParen *> expressionP <* tok TokRParen
 
-objTypeBody :: BodyType F (Commented SourceSpan) -> Pa ObjTypeBody
-objTypeBody bt =
+objTypeBodyP :: BodyType F (Commented SourceSpan) -> Pa ObjTypeBody
+objTypeBodyP bt =
   withMeta $
-    ObjTypeBody bt <$> many (directed objTypeDecl <* tok TokSemi)
+    ObjTypeBody bt <$> many (directedP objTypeDeclP <* tok TokSemi)
 
-objTypeDecl :: Pa ObjTypeDecl
-objTypeDecl =
+objTypeDeclP :: Pa ObjTypeDecl
+objTypeDeclP =
   withMeta $
     ( do
         tok KWAssertPos
-        AssertPosStatement (Witness ()) <$> exprInParen
+        AssertPosStatement (Witness ()) <$> exprInParenP
     )
       <|> ( do
               tok KWReserved
-              ReservedDecl <$> exprInParen
+              ReservedDecl <$> exprInParenP
           )
       <|> ( do
-              bt <- bodyType
-              TypeSubStructure <$> defer body (objTypeBody bt) <*> optionMaybe ident
+              bt <- bodyTypeP
+              TypeSubStructure <$> defer body (objTypeBodyP bt) <*> optionMaybe ident
           )
       <|> ( do
-              mod <- optionMaybe modifier
+              mod <- optionMaybe modifierP
               tok KWReg
               RegisterDecl mod
                 <$> optionMaybe ident
-                <*> exprInParen
-                <*> optionMaybe (tok TokColon *> registerBody)
+                <*> exprInParenP
+                <*> optionMaybe (tok TokColon *> registerBodyP)
           )
 
-modifier :: PaS Modifier
-modifier =
+modifierP :: PaS Modifier
+modifierP =
   withMeta $
     ModifierKeyword
       <$> choice
@@ -240,87 +224,87 @@ modifier =
           tok KWWo >> return Wo
         ]
 
-bitBodyType :: PaS BodyType
-bitBodyType =
+bitBodyTypeP :: PaS BodyType
+bitBodyTypeP =
   withMeta $
     (tok KWStruct >> return Struct)
       <|> (tok KWUnion >> return Union)
 
-bodyType :: PaS BodyType
-bodyType =
+bodyTypeP :: PaS BodyType
+bodyTypeP =
   withMeta $
     (tok KWStruct >> return Struct) <|> (tok KWUnion >> return Union)
 
-registerBody :: Pa RegisterBody
-registerBody = withMeta $ RegisterBody <$> bitBodyType <*> defer body deferredRegisterBody
+registerBodyP :: Pa RegisterBody
+registerBodyP = withMeta $ RegisterBody <$> bitBodyTypeP <*> defer body deferredRegisterBodyP
 
-deferredRegisterBody :: Pa DeferredRegisterBody
-deferredRegisterBody =
+deferredRegisterBodyP :: Pa DeferredRegisterBody
+deferredRegisterBodyP =
   withMeta $
-    DeferredRegisterBody <$> many (directed registerBitsDecl <* tok TokSemi)
+    DeferredRegisterBody <$> many (directedP registerBitsDeclP <* tok TokSemi)
 
-registerBitsDecl :: Pa RegisterBitsDecl
-registerBitsDecl =
+registerBitsDeclP :: Pa RegisterBitsDecl
+registerBitsDeclP =
   withMeta $
     ( do
-        tok KWReserved >> ReservedBits <$> exprInParen
+        tok KWReserved >> ReservedBits <$> exprInParenP
     )
-      <|> (BitsSubStructure <$> registerBody <*> optionMaybe ident)
+      <|> (BitsSubStructure <$> registerBodyP <*> optionMaybe ident)
       <|> ( DefinedBits
-              <$> optionMaybe modifier
+              <$> optionMaybe modifierP
               <*> ident
-              <*> (tok TokColon >> registerBitsTypeRef)
+              <*> (tok TokColon >> registerBitsTypeRefP)
           )
 
-registerBitsTypeRef :: Pa RegisterBitsTypeRef
-registerBitsTypeRef = do
+registerBitsTypeRefP :: Pa RegisterBitsTypeRef
+registerBitsTypeRefP = do
   base <- baseTypeRef
-  recur' <- recur
+  recur' <- recurP
   return (recur' base)
   where
-    recur :: P (RegisterBitsTypeRef Stage1 F A -> RegisterBitsTypeRef Stage1 F A)
-    recur =
+    recurP :: P (RegisterBitsTypeRef Stage1 F A -> RegisterBitsTypeRef Stage1 F A)
+    recurP =
       ( do
           withMeta $ do
-            expr <- tok TokLBracket *> expression <* tok TokRBracket
-            recur' <- recur
+            expr <- tok TokLBracket *> expressionP <* tok TokRBracket
+            recur' <- recurP
             return (\met base -> recur' (RegisterBitsArray base expr met))
       )
         <|> return id
 
     baseTypeRef =
       withMeta $
-        (RegisterBitsJustBits <$> exprInParen)
-          <|> (RegisterBitsAnonymousType (Witness ()) <$> anonymousBitsType)
+        (RegisterBitsJustBits <$> exprInParenP)
+          <|> (RegisterBitsAnonymousType (Witness ()) <$> anonymousBitsTypeP)
           <|> (RegisterBitsReference <$> name)
 
-anonymousBitsType :: Pa AnonymousBitsType
-anonymousBitsType = withMeta $ do
+anonymousBitsTypeP :: Pa AnonymousBitsType
+anonymousBitsTypeP = withMeta $ do
   tok KWEnum
-  AnonymousEnumBody <$> exprInParen <*> defer body enumBody
+  AnonymousEnumBody <$> exprInParenP <*> defer body enumBodyP
 
-bitType :: Pa BitType
-bitType = withMeta $ rawBits <|> enumType
+bitTypeP :: Pa BitType
+bitTypeP = withMeta $ rawBits <|> enumType
   where
-    rawBits = RawBits <$> (tok TokLParen *> expression <* tok TokRParen)
+    rawBits = RawBits <$> (tok TokLParen *> expressionP <* tok TokRParen)
     enumType = do
       tok KWEnum
-      expr <- exprInParen
-      EnumBitType expr <$> defer body enumBody
+      expr <- exprInParenP
+      EnumBitType expr <$> defer body enumBodyP
 
-enumBody :: Pa EnumBody
-enumBody =
+enumBodyP :: Pa EnumBody
+enumBodyP =
   withMeta $
-    EnumBody <$> many (directed enumConstantDecl <* tok TokComma)
+    EnumBody <$> many (directedP enumConstantDeclP <* tok TokComma)
 
-enumConstantDecl :: Pa EnumConstantDecl
-enumConstantDecl =
+enumConstantDeclP :: Pa EnumConstantDecl
+enumConstantDeclP =
   withMeta $
-    (tok KWReserved >> EnumConstantReserved <$> (tok TokEq >> expression))
-      <|> (EnumConstantDecl <$> ident <*> (tok TokEq >> expression))
+    (tok KWReserved >> EnumConstantReserved <$> (tok TokEq >> expressionP))
+      <|> (EnumConstantDecl <$> ident <*> (tok TokEq >> expressionP))
 
-expression :: Pa Expression
-expression = withMeta $
+expressionP :: Pa Expression
+expressionP = withMeta $
   token $ \case
     (TokLitNum num) -> Just (LitNum num)
     (TokIdent i) -> Just $
@@ -373,13 +357,13 @@ defer p0 pb = do
     (sourceName sourcePos)
     <$> p0
 
-packageBody :: Pa PackageBody
-packageBody =
+packageBodyP :: Pa PackageBody
+packageBodyP =
   withMeta $
     PackageBody
       <$> many
-        ( directed $
-            fiddleDecl
+        ( directedP $
+            fiddleDeclP
               <* ( tok TokSemi <|> fail "Expected ';'"
                  )
         )
@@ -409,7 +393,7 @@ name = withMeta $ do
 -- and after and sets the positions and adds it to the annotation.
 withMeta :: P (Commented SourceSpan -> b) -> P b
 withMeta p = do
-  comments <- many comment
+  comments <- many commentP
   start <- getPosition
   fn <- p
   end <- getPosition
@@ -431,7 +415,7 @@ tokKeepComment t' = do
 
 tok :: T -> P (Token SourceSpan)
 tok t' = do
-  many comment
+  many commentP
   Text.Parsec.token
     (\(Token t _) -> show t)
     (\(Token _ (SourceSpan s1 _)) -> s1)
diff --git a/src/Main.hs b/src/Main.hs
index f92d6c6..6fba502 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -1,17 +1,21 @@
 module Main where
 
 import Control.Monad (forM_)
+import Control.Monad.Identity (Identity)
 import Control.Monad.Writer
-import Data.Aeson (encode)
+import Data.Aeson (Value (Null), encode)
 import qualified Data.ByteString.Lazy.Char8 as BL
+import Data.Data (cast)
+import qualified Data.Text as Text
 import qualified Data.Text.IO
+import Data.Typeable
 import GHC.IO.Exception (ExitCode (ExitFailure, ExitSuccess))
 import Language.Fiddle.Ast
 import Language.Fiddle.Compiler (coloredFormat, compile_, printDiagnostic)
 import Language.Fiddle.Compiler.Stage0
 import Language.Fiddle.Compiler.Stage1
 import Language.Fiddle.Compiler.Stage2
-import Language.Fiddle.GenericTree (ToGenericSyntaxTree (toGenericSyntaxTree))
+import Language.Fiddle.GenericTree (GenericSyntaxTree (..), ToGenericSyntaxTree (toGenericSyntaxTree), alterGenericSyntaxTree)
 import qualified Language.Fiddle.Parser
 import qualified Language.Fiddle.Tokenizer
 import qualified System.Environment as System
@@ -25,10 +29,17 @@ main = do
     [filePath] -> do
       text <- Data.Text.IO.readFile filePath
       let (diags, ma) = compile_ $ toStage3 =<< toStage2 =<< toStage1 =<< toStage0 filePath text
-      ec <- 
+      ec <-
         case ma of
           Just ast -> do
-            putStrLn $ BL.unpack $ encode $ toGenericSyntaxTree ast
+            putStrLn $
+              BL.unpack $
+                encode $
+                  alterGenericSyntaxTree cleanupIdentifiers $
+                    toGenericSyntaxTree $
+                      fmap
+                        (const (Nothing :: Maybe Value))
+                        ast
             return ExitSuccess
           Nothing -> do
             putStrLn "\x1b[1;31mCompilation Failed\x1b[0m"
@@ -39,3 +50,15 @@ main = do
     _ -> do
       putStrLn "Wrong Args"
       exitWith (ExitFailure 2)
+
+cleanupIdentifiers :: GenericSyntaxTree Identity a -> Maybe (GenericSyntaxTree Identity a)
+cleanupIdentifiers (SyntaxTreeObject _ _ _ tr)
+  | (Just (Identifier n _)) <- castT tr =
+      Just $ SyntaxTreeValue (Text.unpack n)
+  where
+    castT ::
+      (Typeable t, Typeable f, Typeable a, Typeable t') =>
+      t f a ->
+      Maybe (t' f a)
+    castT = cast
+cleanupIdentifiers _ = Nothing