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.

1 files changed, 276 insertions, 305 deletions

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'