diff --git a/CHANGELOG.md b/CHANGELOG.md
index fd8a424..320f5bf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,12 @@
 
 ## 1.1.0.0
 
+* Change the order of fields in `GroupEntry`; the extension field is now the last field
+* Add `IndexMappable` to help with traversing `CDDL` trees
+* Add an index type parameter to all `CDDL` terms
+* Remove `Codec.CBOR.Cuddle.CDDL.Prelude`
+* Replace `cddlPrelude` with `cddlPostlude`, `prependPrelude` with `appendPostlude`
+* Move `PTerm` to `Codec.CBOR.Cuddle.CDDL.CTree`
 * Remove `CTreeRoot'`
 * Changed the type in `CTreeRoot` to a map of resolved `CTree`s
 * Changed the type of the first argument for `generateCBORTerm` and 
diff --git a/bin/Main.hs b/bin/Main.hs
index 4564d36..03d1e2b 100644
--- a/bin/Main.hs
+++ b/bin/Main.hs
@@ -4,13 +4,14 @@ module Main (main) where
 
 import Codec.CBOR.Cuddle.CBOR.Gen (generateCBORTerm)
 import Codec.CBOR.Cuddle.CBOR.Validator
-import Codec.CBOR.Cuddle.CDDL (Name (..), sortCDDL)
-import Codec.CBOR.Cuddle.CDDL.Prelude (prependPrelude)
+import Codec.CBOR.Cuddle.CDDL (Name (..), fromRules, sortCDDL)
+import Codec.CBOR.Cuddle.CDDL.Postlude (appendPostlude)
 import Codec.CBOR.Cuddle.CDDL.Resolve (
   fullResolveCDDL,
  )
+import Codec.CBOR.Cuddle.IndexMappable (IndexMappable (..))
 import Codec.CBOR.Cuddle.Parser (pCDDL)
-import Codec.CBOR.Cuddle.Pretty ()
+import Codec.CBOR.Cuddle.Pretty (PrettyStage)
 import Codec.CBOR.FlatTerm (toFlatTerm)
 import Codec.CBOR.Pretty (prettyHexEnc)
 import Codec.CBOR.Term (encodeTerm)
@@ -185,26 +186,26 @@ run (Opts cmd cddlFile) = do
         Format fOpts ->
           let
             defs
-              | sort fOpts = sortCDDL res
+              | sort fOpts = fromRules $ sortCDDL res
               | otherwise = res
            in
-            putDocW 80 $ pretty defs
+            putDocW 80 . pretty $ mapIndex @_ @_ @PrettyStage defs
         Validate vOpts ->
           let
             res'
               | vNoPrelude vOpts = res
-              | otherwise = prependPrelude res
+              | otherwise = appendPostlude res
            in
-            case fullResolveCDDL res' of
+            case fullResolveCDDL $ mapIndex res' of
               Left err -> putStrLnErr (show err) >> exitFailure
               Right _ -> exitSuccess
         (GenerateCBOR gOpts) ->
           let
             res'
               | gNoPrelude gOpts = res
-              | otherwise = prependPrelude res
+              | otherwise = appendPostlude res
            in
-            case fullResolveCDDL res' of
+            case fullResolveCDDL $ mapIndex res' of
               Left err -> putStrLnErr (show err) >> exitFailure
               Right mt -> do
                 stdGen <- getStdGen
@@ -220,9 +221,9 @@ run (Opts cmd cddlFile) = do
           let
             res'
               | vcNoPrelude vcOpts = res
-              | otherwise = prependPrelude res
+              | otherwise = res
            in
-            case fullResolveCDDL res' of
+            case fullResolveCDDL $ mapIndex res' of
               Left err -> putStrLnErr (show err) >> exitFailure
               Right mt -> do
                 cbor <- BSC.readFile (vcInput vcOpts)
diff --git a/cuddle.cabal b/cuddle.cabal
index 63e7d0d..ccce127 100644
--- a/cuddle.cabal
+++ b/cuddle.cabal
@@ -49,12 +49,12 @@ library
     Codec.CBOR.Cuddle.CDDL.CTree
     Codec.CBOR.Cuddle.CDDL.CtlOp
     Codec.CBOR.Cuddle.CDDL.Postlude
-    Codec.CBOR.Cuddle.CDDL.Prelude
     Codec.CBOR.Cuddle.CDDL.Resolve
     Codec.CBOR.Cuddle.Comments
     Codec.CBOR.Cuddle.Huddle
     Codec.CBOR.Cuddle.Huddle.HuddleM
     Codec.CBOR.Cuddle.Huddle.Optics
+    Codec.CBOR.Cuddle.IndexMappable
     Codec.CBOR.Cuddle.Parser
     Codec.CBOR.Cuddle.Parser.Lexer
     Codec.CBOR.Cuddle.Pretty
@@ -150,6 +150,7 @@ test-suite cuddle-test
     bytestring,
     cuddle,
     data-default-class,
+    generic-random,
     hspec >=2.11,
     hspec-megaparsec >=2.2,
     megaparsec,
diff --git a/src/Codec/CBOR/Cuddle/CBOR/Gen.hs b/src/Codec/CBOR/Cuddle/CBOR/Gen.hs
index 69a5154..20ae3fc 100644
--- a/src/Codec/CBOR/Cuddle/CBOR/Gen.hs
+++ b/src/Codec/CBOR/Cuddle/CBOR/Gen.hs
@@ -25,10 +25,9 @@ import Codec.CBOR.Cuddle.CDDL (
   Value (..),
   ValueVariant (..),
  )
-import Codec.CBOR.Cuddle.CDDL.CTree (CTree, CTreeRoot (..))
+import Codec.CBOR.Cuddle.CDDL.CTree (CTree, CTreePhase, CTreeRoot (..), PTerm (..))
 import Codec.CBOR.Cuddle.CDDL.CTree qualified as CTree
 import Codec.CBOR.Cuddle.CDDL.CtlOp qualified as CtlOp
-import Codec.CBOR.Cuddle.CDDL.Postlude (PTerm (..))
 import Codec.CBOR.Cuddle.CDDL.Resolve (MonoRef (..), MonoReferenced)
 import Codec.CBOR.Term (Term (..))
 import Codec.CBOR.Term qualified as CBOR
@@ -380,7 +379,7 @@ resolveRef (MRuleRef n) = do
 -- This will throw an error if the generated item does not correspond to a
 -- single CBOR term (e.g. if the name resolves to a group, which cannot be
 -- generated outside a context).
-genForName :: RandomGen g => Name -> M g Term
+genForName :: RandomGen g => Name CTreePhase -> M g Term
 genForName n = do
   (CTreeRoot cddl) <- ask @"cddl"
   case Map.lookup n cddl of
@@ -434,13 +433,13 @@ genValueVariant (VBool b) = pure $ TBool b
 -- Generator functions
 --------------------------------------------------------------------------------
 
-generateCBORTerm :: RandomGen g => CTreeRoot MonoReferenced -> Name -> g -> Term
+generateCBORTerm :: RandomGen g => CTreeRoot MonoReferenced -> Name CTreePhase -> g -> Term
 generateCBORTerm cddl n stdGen =
   let genEnv = GenEnv {cddl}
       genState = GenState {randomSeed = stdGen, depth = 1}
    in evalGen (genForName n) genEnv genState
 
-generateCBORTerm' :: RandomGen g => CTreeRoot MonoReferenced -> Name -> g -> (Term, g)
+generateCBORTerm' :: RandomGen g => CTreeRoot MonoReferenced -> Name CTreePhase -> g -> (Term, g)
 generateCBORTerm' cddl n stdGen =
   let genEnv = GenEnv {cddl}
       genState = GenState {randomSeed = stdGen, depth = 1}
diff --git a/src/Codec/CBOR/Cuddle/CBOR/Validator.hs b/src/Codec/CBOR/Cuddle/CBOR/Validator.hs
index d3894fb..e763078 100644
--- a/src/Codec/CBOR/Cuddle/CBOR/Validator.hs
+++ b/src/Codec/CBOR/Cuddle/CBOR/Validator.hs
@@ -10,7 +10,6 @@ module Codec.CBOR.Cuddle.CBOR.Validator (
 import Codec.CBOR.Cuddle.CDDL hiding (CDDL, Group, Rule)
 import Codec.CBOR.Cuddle.CDDL.CTree
 import Codec.CBOR.Cuddle.CDDL.CtlOp
-import Codec.CBOR.Cuddle.CDDL.Postlude
 import Codec.CBOR.Cuddle.CDDL.Resolve
 import Codec.CBOR.Read
 import Codec.CBOR.Term
@@ -113,7 +112,7 @@ data AMatchedItem = AMatchedItem
 --------------------------------------------------------------------------------
 -- Main entry point
 
-validateCBOR :: BS.ByteString -> Name -> CDDL -> IO ()
+validateCBOR :: BS.ByteString -> Name CTreePhase -> CDDL -> IO ()
 validateCBOR bs rule cddl =
   ( case validateCBOR' bs rule cddl of
       ok@(CBORTermResult _ (Valid _)) -> do
@@ -130,7 +129,7 @@ validateCBOR bs rule cddl =
             )
 
 validateCBOR' ::
-  BS.ByteString -> Name -> CDDL -> CBORTermResult
+  BS.ByteString -> Name CTreePhase -> CDDL -> CBORTermResult
 validateCBOR' bs rule cddl@(CTreeRoot tree) =
   case deserialiseFromBytes decodeTerm (BSL.fromStrict bs) of
     Left e -> error $ show e
diff --git a/src/Codec/CBOR/Cuddle/CDDL.hs b/src/Codec/CBOR/Cuddle/CDDL.hs
index 173dc23..a06ab8a 100644
--- a/src/Codec/CBOR/Cuddle/CDDL.hs
+++ b/src/Codec/CBOR/Cuddle/CDDL.hs
@@ -1,5 +1,9 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
 {-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE OverloadedLabels #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE UndecidableInstances #-}
 
 -- | This module defined the data structure of CDDL as specified in
 --   https://datatracker.ietf.org/doc/rfc8610/
@@ -7,9 +11,9 @@ module Codec.CBOR.Cuddle.CDDL (
   CDDL (..),
   sortCDDL,
   cddlTopLevel,
-  cddlRules,
-  fromRules,
   fromRule,
+  fromRules,
+  appendRules,
   TopLevel (..),
   Name (..),
   Rule (..),
@@ -33,66 +37,95 @@ module Codec.CBOR.Cuddle.CDDL (
   GrpChoice (..),
   unwrap,
   compareRuleName,
+  -- Extension
+  ForAllExtensions,
+  XXTopLevel,
+  XXType2,
+  XCddl,
+  XTerm,
 ) where
 
 import Codec.CBOR.Cuddle.CDDL.CtlOp (CtlOp)
 import Codec.CBOR.Cuddle.Comments (CollectComments (..), Comment, HasComment (..))
 import Data.ByteString qualified as B
 import Data.Default.Class (Default (..))
-import Data.Function (on, (&))
+import Data.Function (on)
 import Data.Hashable (Hashable)
-import Data.List.NonEmpty (NonEmpty (..))
+import Data.List.NonEmpty (NonEmpty (..), prependList)
 import Data.List.NonEmpty qualified as NE
+import Data.Maybe (mapMaybe)
 import Data.String (IsString (..))
 import Data.Text qualified as T
 import Data.TreeDiff (ToExpr)
 import Data.Word (Word64, Word8)
+import GHC.Base (Constraint, Type)
 import GHC.Generics (Generic)
-import Optics.Core ((%), (.~))
-import Optics.Getter (view)
-import Optics.Lens (lens)
+import Optics.Core ((%), (%~), (&))
+
+data family XXTopLevel i
+
+data family XCddl i
+
+data family XTerm i
+
+data family XXType2 i
+
+type ForAllExtensions i (c :: Type -> Constraint) =
+  ( c (XCddl i)
+  , c (XXTopLevel i)
+  , c (XTerm i)
+  , c (XXType2 i)
+  )
 
 -- | The CDDL constructor takes three arguments:
 --     1. Top level comments that precede the first definition
 --     2. The root definition
 --     3. All the other top level comments and definitions
 --   This ensures that `CDDL` is correct by construction.
-data CDDL = CDDL [Comment] Rule [TopLevel]
-  deriving (Eq, Generic, Show, ToExpr)
+data CDDL i = CDDL
+  { rootDefinition :: Rule i
+  , topLevelDefinitions :: [TopLevel i]
+  , cddlExt :: [XXTopLevel i]
+  -- ^ This extension is used for comments that appear before the root definition
+  }
+  deriving (Generic)
+
+deriving instance ForAllExtensions i Eq => Eq (CDDL i)
+
+deriving instance ForAllExtensions i Show => Show (CDDL i)
+
+deriving instance ForAllExtensions i ToExpr => ToExpr (CDDL i)
+
+ruleTopLevel :: TopLevel i -> Maybe (Rule i)
+ruleTopLevel (TopLevelRule r) = Just r
+ruleTopLevel _ = Nothing
 
 -- | Sort the CDDL Rules on the basis of their names
--- Top level comments will be removed!
-sortCDDL :: CDDL -> CDDL
-sortCDDL = fromRules . NE.sortBy (compare `on` ruleName) . cddlRules
-
-cddlTopLevel :: CDDL -> NonEmpty TopLevel
-cddlTopLevel (CDDL cmts cHead cTail) =
-  prependList (TopLevelComment <$> cmts) $ TopLevelRule cHead :| cTail
-  where
-    prependList [] l = l
-    prependList (x : xs) (y :| ys) = prependList xs $ x :| (y : ys)
-
-cddlRules :: CDDL -> NonEmpty Rule
-cddlRules (CDDL _ x tls) = x :| concatMap getRule tls
-  where
-    getRule (TopLevelRule r) = [r]
-    getRule _ = mempty
-
-fromRules :: NonEmpty Rule -> CDDL
-fromRules (x :| xs) = CDDL [] x $ TopLevelRule <$> xs
-
-fromRule :: Rule -> CDDL
-fromRule x = CDDL [] x []
-
-instance Semigroup CDDL where
-  CDDL aComments aHead aTail <> CDDL bComments bHead bTail =
-    CDDL aComments aHead $
-      aTail <> fmap TopLevelComment bComments <> (TopLevelRule bHead : bTail)
-
-data TopLevel
-  = TopLevelRule Rule
-  | TopLevelComment Comment
-  deriving (Eq, Generic, Show, ToExpr)
+sortCDDL :: CDDL i -> NonEmpty (Rule i)
+sortCDDL (CDDL r rs _) = NE.sortBy (compare `on` name . ruleName) $ r :| mapMaybe ruleTopLevel rs
+
+fromRules :: Monoid (XCddl i) => NonEmpty (Rule i) -> CDDL i
+fromRules (x :| xs) = CDDL x (TopLevelRule <$> xs) mempty
+
+fromRule :: Monoid (XCddl i) => Rule i -> CDDL i
+fromRule x = CDDL x [] mempty
+
+appendRules :: CDDL i -> [Rule i] -> CDDL i
+appendRules cddl rs = cddl & #topLevelDefinitions %~ (<> fmap TopLevelRule rs)
+
+cddlTopLevel :: CDDL i -> NonEmpty (TopLevel i)
+cddlTopLevel (CDDL r tls e) = prependList (XXTopLevel <$> e) $ TopLevelRule r :| tls
+
+data TopLevel i
+  = TopLevelRule (Rule i)
+  | XXTopLevel (XXTopLevel i)
+  deriving (Generic)
+
+deriving instance ForAllExtensions i Eq => Eq (TopLevel i)
+
+deriving instance ForAllExtensions i Show => Show (TopLevel i)
+
+deriving instance ForAllExtensions i ToExpr => ToExpr (TopLevel i)
 
 -- |
 --  A name can consist of any of the characters from the set {"A" to
@@ -117,23 +150,30 @@ data TopLevel
 --
 --  *  Rule names (types or groups) do not appear in the actual CBOR
 --      encoding, but names used as "barewords" in member keys do.
-data Name = Name
+data Name i = Name
   { name :: T.Text
-  , nameComment :: Comment
+  , nameExt :: XTerm i
   }
-  deriving (Eq, Generic, Ord, Show)
-  deriving anyclass (ToExpr)
+  deriving (Generic)
+
+deriving instance Eq (XTerm i) => Eq (Name i)
+
+deriving instance Ord (XTerm i) => Ord (Name i)
+
+deriving instance Show (XTerm i) => Show (Name i)
+
+deriving instance ToExpr (XTerm i) => ToExpr (Name i)
 
-instance IsString Name where
+instance Monoid (XTerm i) => IsString (Name i) where
   fromString x = Name (T.pack x) mempty
 
-instance HasComment Name where
-  commentL = lens nameComment (\x y -> x {nameComment = y})
+instance HasComment (XTerm i) => HasComment (Name i) where
+  commentL = #nameExt % commentL
 
-instance CollectComments Name where
-  collectComments (Name _ c) = [c]
+instance CollectComments (XTerm i) => CollectComments (Name i) where
+  collectComments (Name _ c) = collectComments c
 
-instance Hashable Name
+instance Hashable (XTerm i) => Hashable (Name i)
 
 -- |
 --   assignt = "=" / "/="
@@ -168,17 +208,27 @@ data Assign = AssignEq | AssignExt
 --
 --   Generic rules can be used for establishing names for both types and
 --   groups.
-newtype GenericParam = GenericParam (NE.NonEmpty Name)
-  deriving (Eq, Generic, Show)
+newtype GenericParam i = GenericParam (NE.NonEmpty (Name i))
+  deriving (Generic)
   deriving newtype (Semigroup)
-  deriving anyclass (ToExpr)
 
-newtype GenericArg = GenericArg (NE.NonEmpty Type1)
-  deriving (Eq, Generic, Show)
+deriving instance Eq (XTerm i) => Eq (GenericParam i)
+
+deriving instance Show (XTerm i) => Show (GenericParam i)
+
+deriving anyclass instance ToExpr (XTerm i) => ToExpr (GenericParam i)
+
+newtype GenericArg i = GenericArg (NE.NonEmpty (Type1 i))
+  deriving (Generic)
   deriving newtype (Semigroup)
-  deriving anyclass (ToExpr)
 
-instance CollectComments GenericArg
+deriving instance ForAllExtensions i Eq => Eq (GenericArg i)
+
+deriving instance ForAllExtensions i Show => Show (GenericArg i)
+
+deriving anyclass instance ForAllExtensions i ToExpr => ToExpr (GenericArg i)
+
+instance ForAllExtensions i CollectComments => CollectComments (GenericArg i)
 
 -- |
 --  rule = typename [genericparm] S assignt S type
@@ -203,20 +253,25 @@ instance CollectComments GenericArg
 --   clear immediately either whether "b" stands for a group or a type --
 --   this semantic processing may need to span several levels of rule
 --   definitions before a determination can be made.)
-data Rule = Rule
-  { ruleName :: Name
-  , ruleGenParam :: Maybe GenericParam
+data Rule i = Rule
+  { ruleName :: Name i
+  , ruleGenParam :: Maybe (GenericParam i)
   , ruleAssign :: Assign
-  , ruleTerm :: TypeOrGroup
-  , ruleComment :: Comment
+  , ruleTerm :: TypeOrGroup i
+  , ruleExt :: XTerm i
   }
-  deriving (Eq, Generic, Show)
-  deriving anyclass (ToExpr)
+  deriving (Generic)
+
+deriving instance ForAllExtensions i Eq => Eq (Rule i)
+
+deriving instance ForAllExtensions i Show => Show (Rule i)
+
+deriving instance ForAllExtensions i ToExpr => ToExpr (Rule i)
 
-instance HasComment Rule where
-  commentL = lens ruleComment (\x y -> x {ruleComment = y})
+instance HasComment (XTerm i) => HasComment (Rule i) where
+  commentL = #ruleExt % commentL
 
-compareRuleName :: Rule -> Rule -> Ordering
+compareRuleName :: Ord (XTerm i) => Rule i -> Rule i -> Ordering
 compareRuleName = compare `on` ruleName
 
 -- |
@@ -235,11 +290,16 @@ data TyOp = RangeOp RangeBound | CtrlOp CtlOp
   deriving (Eq, Generic, Show)
   deriving anyclass (ToExpr)
 
-data TypeOrGroup = TOGType Type0 | TOGGroup GroupEntry
-  deriving (Eq, Generic, Show)
-  deriving anyclass (ToExpr)
+data TypeOrGroup i = TOGType (Type0 i) | TOGGroup (GroupEntry i)
+  deriving (Generic)
+
+deriving instance ForAllExtensions i Eq => Eq (TypeOrGroup i)
+
+deriving instance ForAllExtensions i Show => Show (TypeOrGroup i)
+
+deriving instance ForAllExtensions i ToExpr => ToExpr (TypeOrGroup i)
 
-instance CollectComments TypeOrGroup
+instance ForAllExtensions i CollectComments => CollectComments (TypeOrGroup i)
 
 {-- |
    The group that is used to define a map or an array can often be reused in the
@@ -290,7 +350,7 @@ instance CollectComments TypeOrGroup
    described as "threading in" the group or type inside the referenced type,
    which suggested the thread-like "~" character.)
 -}
-unwrap :: TypeOrGroup -> Maybe Group
+unwrap :: TypeOrGroup i -> Maybe (Group i)
 unwrap (TOGType (Type0 (Type1 t2 Nothing _ NE.:| []))) = case t2 of
   T2Map g -> Just g
   T2Array g -> Just g
@@ -301,70 +361,83 @@ unwrap _ = Nothing
 -- A type can be given as a choice between one or more types.  The
 --   choice matches a data item if the data item matches any one of the
 --   types given in the choice.
-newtype Type0 = Type0 {t0Type1 :: NE.NonEmpty Type1}
-  deriving (Eq, Generic, Show)
+newtype Type0 i = Type0 {t0Type1 :: NE.NonEmpty (Type1 i)}
+  deriving (Generic)
   deriving newtype (Semigroup)
-  deriving anyclass (ToExpr)
 
-instance HasComment Type0 where
-  commentL = lens (view commentL . t0Type1) (\(Type0 x) y -> Type0 $ x & commentL .~ y)
+deriving instance ForAllExtensions i Eq => Eq (Type0 i)
+
+deriving instance ForAllExtensions i Show => Show (Type0 i)
 
-instance CollectComments Type0
+deriving anyclass instance ForAllExtensions i ToExpr => ToExpr (Type0 i)
+
+instance ForAllExtensions i CollectComments => CollectComments (Type0 i)
 
 -- |
 -- Two types can be combined with a range operator (see below)
-data Type1 = Type1
-  { t1Main :: Type2
-  , t1TyOp :: Maybe (TyOp, Type2)
-  , t1Comment :: Comment
+data Type1 i = Type1
+  { t1Main :: Type2 i
+  , t1TyOp :: Maybe (TyOp, Type2 i)
+  , t1Comment :: XTerm i
   }
-  deriving (Eq, Generic, Show)
-  deriving anyclass (ToExpr, Default)
+  deriving (Generic)
 
-instance HasComment Type1 where
-  commentL = lens t1Comment (\x y -> x {t1Comment = y})
+deriving instance ForAllExtensions i Eq => Eq (Type1 i)
 
-instance CollectComments Type1 where
-  collectComments (Type1 m tyOp c) = c : collectComments m <> collectComments (fmap snd tyOp)
+deriving instance ForAllExtensions i Show => Show (Type1 i)
 
-data Type2
+deriving instance ForAllExtensions i ToExpr => ToExpr (Type1 i)
+
+instance HasComment (XTerm i) => HasComment (Type1 i) where
+  commentL = #t1Comment % commentL
+
+instance ForAllExtensions i CollectComments => CollectComments (Type1 i) where
+  collectComments (Type1 m tyOp c) = collectComments c <> collectComments m <> collectComments (fmap snd tyOp)
+
+data Type2 i
   = -- | A type can be just a single value (such as 1 or "icecream" or
     --   h'0815'), which matches only a data item with that specific value
     --   (no conversions defined),
     T2Value Value
   | -- | or be defined by a rule giving a meaning to a name (possibly after
     --   supplying generic arguments as required by the generic parameters)
-    T2Name Name (Maybe GenericArg)
+    T2Name (Name i) (Maybe (GenericArg i))
   | -- | or be defined in a parenthesized type expression (parentheses may be
     --   necessary to override some operator precedence),
-    T2Group Type0
+    T2Group (Type0 i)
   | -- | a map expression, which matches a valid CBOR map the key/value pairs
     --  of which can be ordered in such a way that the resulting sequence
     --  matches the group expression, or
-    T2Map Group
+    T2Map (Group i)
   | -- | an array expression, which matches a CBOR array the elements of which
     -- when taken as values and complemented by a wildcard (matches
     -- anything) key each -- match the group, or
-    T2Array Group
+    T2Array (Group i)
   | -- | an "unwrapped" group (see Section 3.7), which matches the group
     --  inside a type defined as a map or an array by wrapping the group, or
-    T2Unwrapped Name (Maybe GenericArg)
+    T2Unwrapped (Name i) (Maybe (GenericArg i))
   | -- | an enumeration expression, which matches any value that is within the
     --  set of values that the values of the group given can take, or
-    T2Enum Group
-  | T2EnumRef Name (Maybe GenericArg)
+    T2Enum (Group i)
+  | T2EnumRef (Name i) (Maybe (GenericArg i))
   | -- | a tagged data item, tagged with the "uint" given and containing the
     --  type given as the tagged value, or
-    T2Tag (Maybe Word64) Type0
+    T2Tag (Maybe Word64) (Type0 i)
   | -- | a data item of a major type (given by the DIGIT), optionally
     --  constrained to the additional information given by the uint, or
     T2DataItem Word8 (Maybe Word64)
   | -- | Any data item
     T2Any
-  deriving (Eq, Generic, Show, Default)
-  deriving anyclass (ToExpr)
+  | XXType2 (XXType2 i)
+  deriving (Generic)
 
-instance CollectComments Type2
+deriving instance ForAllExtensions i Eq => Eq (Type2 i)
+
+deriving instance ForAllExtensions i Show => Show (Type2 i)
+
+deriving instance ForAllExtensions i ToExpr => ToExpr (Type2 i)
+
+instance ForAllExtensions i CollectComments => CollectComments (Type2 i)
 
 -- |
 --  An optional _occurrence_ indicator can be given in front of a group
@@ -393,29 +466,39 @@ instance Hashable OccurrenceIndicator
 -- |
 --   A group matches any sequence of key/value pairs that matches any of
 --   the choices given (again using PEG semantics).
-newtype Group = Group {unGroup :: NE.NonEmpty GrpChoice}
-  deriving (Eq, Generic, Show)
+newtype Group i = Group {unGroup :: NE.NonEmpty (GrpChoice i)}
+  deriving (Generic)
   deriving newtype (Semigroup)
-  deriving anyclass (ToExpr)
 
-instance HasComment Group where
-  commentL = lens unGroup (\x y -> x {unGroup = y}) % commentL
+deriving instance ForAllExtensions i Eq => Eq (Group i)
 
-instance CollectComments Group where
+deriving instance ForAllExtensions i Show => Show (Group i)
+
+deriving anyclass instance ForAllExtensions i ToExpr => ToExpr (Group i)
+
+instance HasComment (XTerm i) => HasComment (Group i) where
+  commentL = #unGroup % commentL
+
+instance ForAllExtensions i CollectComments => CollectComments (Group i) where
   collectComments (Group xs) = concatMap collectComments xs
 
-data GrpChoice = GrpChoice
-  { gcGroupEntries :: [GroupEntry]
-  , gcComment :: Comment
+data GrpChoice i = GrpChoice
+  { gcGroupEntries :: [GroupEntry i]
+  , gcComment :: XTerm i
   }
-  deriving (Eq, Generic, Show)
-  deriving anyclass (ToExpr)
+  deriving (Generic)
 
-instance HasComment GrpChoice where
-  commentL = lens gcComment (\x y -> x {gcComment = y})
+deriving instance ForAllExtensions i Eq => Eq (GrpChoice i)
 
-instance CollectComments GrpChoice where
-  collectComments (GrpChoice ges c) = c : concatMap collectComments ges
+deriving instance ForAllExtensions i Show => Show (GrpChoice i)
+
+deriving instance ForAllExtensions i ToExpr => ToExpr (GrpChoice i)
+
+instance HasComment (XTerm i) => HasComment (GrpChoice i) where
+  commentL = #gcComment % commentL
+
+instance ForAllExtensions i CollectComments => CollectComments (GrpChoice i) where
+  collectComments (GrpChoice ges c) = collectComments c <> concatMap collectComments ges
 
 -- |
 --  A group entry can be given by a value type, which needs to be matched
@@ -424,26 +507,38 @@ instance CollectComments GrpChoice where
 --  the memberkey is given.  If the memberkey is not given, the entry can
 --  only be used for matching arrays, not for maps.  (See below for how
 --  that is modified by the occurrence indicator.)
-data GroupEntry = GroupEntry
+data GroupEntry i = GroupEntry
   { geOccurrenceIndicator :: Maybe OccurrenceIndicator
-  , geComment :: Comment
-  , geVariant :: GroupEntryVariant
+  , geVariant :: GroupEntryVariant i
+  , geExt :: XTerm i
   }
-  deriving (Eq, Show, Generic, ToExpr)
+  deriving (Generic)
+
+deriving instance ForAllExtensions i Eq => Eq (GroupEntry i)
+
+deriving instance ForAllExtensions i Show => Show (GroupEntry i)
 
-instance CollectComments GroupEntry where
-  collectComments (GroupEntry _ c x) = c : collectComments x
+deriving instance ForAllExtensions i ToExpr => ToExpr (GroupEntry i)
 
-data GroupEntryVariant
-  = GEType (Maybe MemberKey) Type0
-  | GERef Name (Maybe GenericArg)
-  | GEGroup Group
-  deriving (Eq, Show, Generic, ToExpr)
+instance ForAllExtensions i CollectComments => CollectComments (GroupEntry i) where
+  collectComments (GroupEntry _ c x) = collectComments c <> collectComments x
 
-instance HasComment GroupEntry where
-  commentL = lens geComment (\x y -> x {geComment = y})
+data GroupEntryVariant i
+  = GEType (Maybe (MemberKey i)) (Type0 i)
+  | GERef (Name i) (Maybe (GenericArg i))
+  | GEGroup (Group i)
+  deriving (Generic)
 
-instance CollectComments GroupEntryVariant where
+deriving instance ForAllExtensions i Eq => Eq (GroupEntryVariant i)
+
+deriving instance ForAllExtensions i Show => Show (GroupEntryVariant i)
+
+deriving instance ForAllExtensions i ToExpr => ToExpr (GroupEntryVariant i)
+
+instance HasComment (XTerm i) => HasComment (GroupEntry i) where
+  commentL = #geExt % commentL
+
+instance ForAllExtensions i CollectComments => CollectComments (GroupEntryVariant i) where
   collectComments (GEType _ t0) = collectComments t0
   collectComments (GERef n mga) = collectComments n <> collectComments mga
   collectComments (GEGroup g) = collectComments g
@@ -456,12 +551,17 @@ instance CollectComments GroupEntryVariant where
 --  member of the key type, unless a cut preceding it in the group
 --  applies (see Section 3.5.4 for how map matching is influenced by the
 --  presence of the cuts denoted by "^" or ":" in previous entries).
-data MemberKey
-  = MKType Type1
-  | MKBareword Name
+data MemberKey i
+  = MKType (Type1 i)
+  | MKBareword (Name i)
   | MKValue Value
-  deriving (Eq, Generic, Show)
-  deriving anyclass (ToExpr)
+  deriving (Generic)
+
+deriving instance ForAllExtensions i Eq => Eq (MemberKey i)
+
+deriving instance ForAllExtensions i Show => Show (MemberKey i)
+
+deriving instance ForAllExtensions i ToExpr => ToExpr (MemberKey i)
 
 data Value = Value ValueVariant Comment
   deriving (Eq, Generic, Show, Default)
diff --git a/src/Codec/CBOR/Cuddle/CDDL/CTree.hs b/src/Codec/CBOR/Cuddle/CDDL/CTree.hs
index 26963c6..4fda00d 100644
--- a/src/Codec/CBOR/Cuddle/CDDL/CTree.hs
+++ b/src/Codec/CBOR/Cuddle/CDDL/CTree.hs
@@ -9,11 +9,17 @@ import Codec.CBOR.Cuddle.CDDL (
   OccurrenceIndicator,
   RangeBound,
   Value,
+  XCddl,
+  XTerm,
+  XXTopLevel,
+  XXType2,
  )
 import Codec.CBOR.Cuddle.CDDL.CtlOp
-import Codec.CBOR.Cuddle.CDDL.Postlude (PTerm)
+import Codec.CBOR.Cuddle.Comments (Comment)
+import Data.Hashable (Hashable)
 import Data.List.NonEmpty qualified as NE
 import Data.Map.Strict qualified as Map
+import Data.Void (Void)
 import Data.Word (Word64)
 import GHC.Generics (Generic)
 
@@ -29,6 +35,20 @@ import GHC.Generics (Generic)
 
 type family CTreeExt i
 
+data CTreePhase
+
+newtype instance XTerm CTreePhase = CTreeXTerm Comment
+  deriving (Generic, Show, Eq, Ord, Hashable, Semigroup, Monoid)
+
+newtype instance XXTopLevel CTreePhase = CTreeXXTopLevel Comment
+  deriving (Generic, Show, Eq, Ord, Hashable)
+
+newtype instance XCddl CTreePhase = CTreeXCddl [Comment]
+  deriving (Generic, Show, Eq, Ord, Hashable)
+
+newtype instance XXType2 CTreePhase = CTreeXXType2 Void
+  deriving (Generic, Show, Eq, Ord, Hashable)
+
 data CTree i
   = Literal Value
   | Postlude PTerm
@@ -77,7 +97,58 @@ traverseCTree atExt _ (CTreeE x) = atExt x
 
 type Node i = CTreeExt i
 
-newtype CTreeRoot i = CTreeRoot (Map.Map Name (CTree i))
+newtype CTreeRoot i = CTreeRoot (Map.Map (Name CTreePhase) (CTree i))
   deriving (Generic)
 
 deriving instance Show (CTree i) => Show (CTreeRoot i)
+
+-- |
+--
+--  CDDL predefines a number of names.  This subsection summarizes these
+--  names, but please see Appendix D for the exact definitions.
+--
+--  The following keywords for primitive datatypes are defined:
+--
+--  "bool"  Boolean value (major type 7, additional information 20
+--    or 21).
+--
+--  "uint"  An unsigned integer (major type 0).
+--
+--  "nint"  A negative integer (major type 1).
+--
+--  "int"  An unsigned integer or a negative integer.
+--
+--  "float16"  A number representable as a half-precision float [IEEE754]
+--    (major type 7, additional information 25).
+--
+--  "float32"  A number representable as a single-precision float
+--    [IEEE754] (major type 7, additional information 26).
+--
+--
+--  "float64"  A number representable as a double-precision float
+--    [IEEE754] (major type 7, additional information 27).
+--
+--  "float"  One of float16, float32, or float64.
+--
+--  "bstr" or "bytes"  A byte string (major type 2).
+--
+--  "tstr" or "text"  Text string (major type 3).
+--
+--  (Note that there are no predefined names for arrays or maps; these
+--  are defined with the syntax given below.)
+data PTerm
+  = PTBool
+  | PTUInt
+  | PTNInt
+  | PTInt
+  | PTHalf
+  | PTFloat
+  | PTDouble
+  | PTBytes
+  | PTText
+  | PTAny
+  | PTNil
+  | PTUndefined
+  deriving (Eq, Generic, Ord, Show)
+
+instance Hashable PTerm
diff --git a/src/Codec/CBOR/Cuddle/CDDL/Postlude.hs b/src/Codec/CBOR/Cuddle/CDDL/Postlude.hs
index 0f5c0a4..b242283 100644
--- a/src/Codec/CBOR/Cuddle/CDDL/Postlude.hs
+++ b/src/Codec/CBOR/Cuddle/CDDL/Postlude.hs
@@ -1,55 +1,72 @@
+{-# LANGUAGE OverloadedStrings #-}
+
 module Codec.CBOR.Cuddle.CDDL.Postlude where
 
-import Data.Hashable (Hashable)
-import GHC.Generics (Generic)
+import Codec.CBOR.Cuddle.CDDL (CDDL (..), TopLevel (..), XTerm, XXType2, appendRules)
+import Codec.CBOR.Cuddle.IndexMappable (IndexMappable (..))
+import Codec.CBOR.Cuddle.Parser (ParserStage, pCDDL)
+import Data.Maybe (mapMaybe)
+import Text.Megaparsec (errorBundlePretty, parse)
 
--- |
---
---  CDDL predefines a number of names.  This subsection summarizes these
---  names, but please see Appendix D for the exact definitions.
---
---  The following keywords for primitive datatypes are defined:
---
---  "bool"  Boolean value (major type 7, additional information 20
---    or 21).
---
---  "uint"  An unsigned integer (major type 0).
---
---  "nint"  A negative integer (major type 1).
---
---  "int"  An unsigned integer or a negative integer.
---
---  "float16"  A number representable as a half-precision float [IEEE754]
---    (major type 7, additional information 25).
---
---  "float32"  A number representable as a single-precision float
---    [IEEE754] (major type 7, additional information 26).
---
---
---  "float64"  A number representable as a double-precision float
---    [IEEE754] (major type 7, additional information 27).
---
---  "float"  One of float16, float32, or float64.
---
---  "bstr" or "bytes"  A byte string (major type 2).
---
---  "tstr" or "text"  Text string (major type 3).
---
---  (Note that there are no predefined names for arrays or maps; these
---  are defined with the syntax given below.)
-data PTerm
-  = PTBool
-  | PTUInt
-  | PTNInt
-  | PTInt
-  | PTHalf
-  | PTFloat
-  | PTDouble
-  | PTBytes
-  | PTText
-  | PTAny
-  | PTNil
-  | PTUndefined
-  deriving (Eq, Generic, Ord, Show)
+-- TODO switch to quasiquotes
+cddlPostlude :: CDDL ParserStage
+cddlPostlude =
+  either (error . errorBundlePretty) id $
+    parse
+      pCDDL
+      "<HARDCODED>"
+      " any = # \
+      \ uint = #0 \
+      \ nint = #1 \
+      \ int = uint / nint \
+      \ \
+      \ bstr = #2 \
+      \ bytes = bstr \
+      \ tstr = #3 \
+      \ text = tstr \
+      \ \
+      \ tdate = #6.0(tstr) \
+      \ time = #6.1(number) \
+      \ number = int / float \
+      \ biguint = #6.2(bstr) \
+      \ bignint = #6.3(bstr) \
+      \ bigint = biguint / bignint \
+      \ integer = int / bigint \
+      \ unsigned = uint / biguint \
+      \ decfrac = #6.4([e10: int, m: integer]) \
+      \ bigfloat = #6.5([e2: int, m: integer]) \
+      \ eb64url = #6.21(any) \
+      \ eb64legacy = #6.22(any) \
+      \ eb16 = #6.23(any) \
+      \ encoded-cbor = #6.24(bstr) \
+      \ uri = #6.32(tstr) \
+      \ b64url = #6.33(tstr) \
+      \ b64legacy = #6.34(tstr) \
+      \ regexp = #6.35(tstr) \
+      \ mime-message = #6.36(tstr) \
+      \ cbor-any = #6.55799(any) \
+      \ float16 = #7.25 \
+      \ float32 = #7.26 \
+      \ float64 = #7.27 \
+      \ float16-32 = float16 / float32 \
+      \ float32-64 = float32 / float64 \
+      \ float = float16-32 / float64 \
+      \  \
+      \ false = #7.20 \
+      \ true = #7.21 \
+      \ bool = false / true \
+      \ nil = #7.22 \
+      \ null = nil \
+      \ undefined = #7.23"
 
-instance Hashable PTerm
+appendPostlude ::
+  ( IndexMappable XXType2 ParserStage i
+  , IndexMappable XTerm ParserStage i
+  ) =>
+  CDDL i -> CDDL i
+appendPostlude cddl = appendRules cddl $ mapIndex <$> (r : rs)
+  where
+    CDDL r tls _ = cddlPostlude
+    f (TopLevelRule x) = Just x
+    f _ = Nothing
+    rs = mapMaybe f tls
diff --git a/src/Codec/CBOR/Cuddle/CDDL/Prelude.hs b/src/Codec/CBOR/Cuddle/CDDL/Prelude.hs
deleted file mode 100644
index 3bcbe1a..0000000
--- a/src/Codec/CBOR/Cuddle/CDDL/Prelude.hs
+++ /dev/null
@@ -1,61 +0,0 @@
-{-# LANGUAGE OverloadedStrings #-}
-
-module Codec.CBOR.Cuddle.CDDL.Prelude (prependPrelude) where
-
-import Codec.CBOR.Cuddle.CDDL (CDDL (..))
-import Codec.CBOR.Cuddle.Parser (pCDDL)
-import Text.Megaparsec (errorBundlePretty, parse)
-
--- TODO switch to quasiquotes
-cddlPrelude :: CDDL
-cddlPrelude =
-  either (error . errorBundlePretty) id $
-    parse
-      pCDDL
-      "<HARDCODED>"
-      " any = # \
-      \ uint = #0 \
-      \ nint = #1 \
-      \ int = uint / nint \
-      \ \
-      \ bstr = #2 \
-      \ bytes = bstr \
-      \ tstr = #3 \
-      \ text = tstr \
-      \ \
-      \ tdate = #6.0(tstr) \
-      \ time = #6.1(number) \
-      \ number = int / float \
-      \ biguint = #6.2(bstr) \
-      \ bignint = #6.3(bstr) \
-      \ bigint = biguint / bignint \
-      \ integer = int / bigint \
-      \ unsigned = uint / biguint \
-      \ decfrac = #6.4([e10: int, m: integer]) \
-      \ bigfloat = #6.5([e2: int, m: integer]) \
-      \ eb64url = #6.21(any) \
-      \ eb64legacy = #6.22(any) \
-      \ eb16 = #6.23(any) \
-      \ encoded-cbor = #6.24(bstr) \
-      \ uri = #6.32(tstr) \
-      \ b64url = #6.33(tstr) \
-      \ b64legacy = #6.34(tstr) \
-      \ regexp = #6.35(tstr) \
-      \ mime-message = #6.36(tstr) \
-      \ cbor-any = #6.55799(any) \
-      \ float16 = #7.25 \
-      \ float32 = #7.26 \
-      \ float64 = #7.27 \
-      \ float16-32 = float16 / float32 \
-      \ float32-64 = float32 / float64 \
-      \ float = float16-32 / float64 \
-      \  \
-      \ false = #7.20 \
-      \ true = #7.21 \
-      \ bool = false / true \
-      \ nil = #7.22 \
-      \ null = nil \
-      \ undefined = #7.23"
-
-prependPrelude :: CDDL -> CDDL
-prependPrelude = (cddlPrelude <>)
diff --git a/src/Codec/CBOR/Cuddle/CDDL/Resolve.hs b/src/Codec/CBOR/Cuddle/CDDL/Resolve.hs
index bd3e29d..6ca2392 100644
--- a/src/Codec/CBOR/Cuddle/CDDL/Resolve.hs
+++ b/src/Codec/CBOR/Cuddle/CDDL/Resolve.hs
@@ -48,10 +48,12 @@ import Codec.CBOR.Cuddle.CDDL as CDDL
 import Codec.CBOR.Cuddle.CDDL.CTree (
   CTree (..),
   CTreeExt,
+  CTreePhase,
   CTreeRoot (..),
+  PTerm (..),
+  XXType2 (..),
  )
 import Codec.CBOR.Cuddle.CDDL.CTree qualified as CTree
-import Codec.CBOR.Cuddle.CDDL.Postlude (PTerm (..))
 import Control.Monad.Except (ExceptT (..), runExceptT)
 import Control.Monad.Reader (Reader, ReaderT (..), runReader)
 import Control.Monad.State.Strict (StateT (..))
@@ -64,11 +66,12 @@ import Data.List (foldl')
 import Data.List.NonEmpty qualified as NE
 import Data.Map.Strict qualified as Map
 import Data.Text qualified as T
+import Data.Void (absurd)
 import GHC.Generics (Generic)
 import Optics.Core
 
 data ProvidedParameters a = ProvidedParameters
-  { parameters :: [Name]
+  { parameters :: [Name CTreePhase]
   , underlying :: a
   }
   deriving (Generic, Functor, Show, Eq, Foldable, Traversable)
@@ -83,33 +86,33 @@ type instance CTreeExt Parametrised = ProvidedParameters (CTree Parametrised)
 -- 1. Rule extensions
 --------------------------------------------------------------------------------
 
-newtype PartialCTreeRoot i = PartialCTreeRoot (Map.Map Name (ProvidedParameters (CTree i)))
+newtype PartialCTreeRoot i = PartialCTreeRoot (Map.Map (Name CTreePhase) (ProvidedParameters (CTree i)))
   deriving (Generic)
 
-type CDDLMap = Map.Map Name (ProvidedParameters TypeOrGroup)
+type CDDLMap = Map.Map (Name CTreePhase) (ProvidedParameters (TypeOrGroup CTreePhase))
 
-toParametrised :: a -> Maybe GenericParam -> ProvidedParameters a
+toParametrised :: a -> Maybe (GenericParam CTreePhase) -> ProvidedParameters a
 toParametrised a Nothing = ProvidedParameters [] a
 toParametrised a (Just (GenericParam gps)) = ProvidedParameters (NE.toList gps) a
 
-asMap :: CDDL -> CDDLMap
+asMap :: CDDL CTreePhase -> CDDLMap
 asMap cddl = foldl' go Map.empty rules
   where
     rules = cddlTopLevel cddl
-    go x (TopLevelComment _) = x
+    go x (XXTopLevel _) = x
     go x (TopLevelRule r) = assignOrExtend x r
 
-    assignOrExtend :: CDDLMap -> Rule -> CDDLMap
+    assignOrExtend :: CDDLMap -> Rule CTreePhase -> CDDLMap
     assignOrExtend m (Rule n gps assign tog _) = case assign of
       -- Equals assignment
       AssignEq -> Map.insert n (toParametrised tog gps) m
       AssignExt -> Map.alter (extend tog gps) n m
 
     extend ::
-      TypeOrGroup ->
-      Maybe GenericParam ->
-      Maybe (ProvidedParameters TypeOrGroup) ->
-      Maybe (ProvidedParameters TypeOrGroup)
+      TypeOrGroup CTreePhase ->
+      Maybe (GenericParam CTreePhase) ->
+      Maybe (ProvidedParameters (TypeOrGroup CTreePhase)) ->
+      Maybe (ProvidedParameters (TypeOrGroup CTreePhase))
     extend tog _gps (Just existing) = case (underlying existing, tog) of
       (TOGType _, TOGType (Type0 new)) ->
         Just $
@@ -139,7 +142,7 @@ type instance CTreeExt OrReferenced = OrRef
 -- | Indicates that an item may be referenced rather than defined.
 data OrRef
   = -- | Reference to another node with possible generic arguments supplied
-    Ref Name [CTree OrReferenced]
+    Ref (Name CTreePhase) [CTree OrReferenced]
   deriving (Eq, Show)
 
 type RefCTree = PartialCTreeRoot OrReferenced
@@ -155,19 +158,19 @@ buildRefCTree :: CDDLMap -> RefCTree
 buildRefCTree rules = PartialCTreeRoot $ toCTreeRule <$> rules
   where
     toCTreeRule ::
-      ProvidedParameters TypeOrGroup ->
+      ProvidedParameters (TypeOrGroup CTreePhase) ->
       ProvidedParameters (CTree OrReferenced)
     toCTreeRule = fmap toCTreeTOG
 
-    toCTreeTOG :: TypeOrGroup -> CTree OrReferenced
+    toCTreeTOG :: TypeOrGroup CTreePhase -> CTree OrReferenced
     toCTreeTOG (TOGType t0) = toCTreeT0 t0
     toCTreeTOG (TOGGroup ge) = toCTreeGroupEntry ge
 
-    toCTreeT0 :: Type0 -> CTree OrReferenced
+    toCTreeT0 :: Type0 CTreePhase -> CTree OrReferenced
     toCTreeT0 (Type0 (t1 NE.:| [])) = toCTreeT1 t1
     toCTreeT0 (Type0 xs) = CTree.Choice $ toCTreeT1 <$> xs
 
-    toCTreeT1 :: Type1 -> CTree OrReferenced
+    toCTreeT1 :: Type1 CTreePhase -> CTree OrReferenced
     toCTreeT1 (Type1 t2 Nothing _) = toCTreeT2 t2
     toCTreeT1 (Type1 t2 (Just (op, t2')) _) = case op of
       RangeOp bound ->
@@ -183,7 +186,7 @@ buildRefCTree rules = PartialCTreeRoot $ toCTreeRule <$> rules
           , CTree.controller = toCTreeT2 t2'
           }
 
-    toCTreeT2 :: Type2 -> CTree OrReferenced
+    toCTreeT2 :: Type2 CTreePhase -> CTree OrReferenced
     toCTreeT2 (T2Value v) = CTree.Literal v
     toCTreeT2 (T2Name n garg) = CTreeE $ Ref n (fromGenArgs garg)
     toCTreeT2 (T2Group t0) =
@@ -208,6 +211,7 @@ buildRefCTree rules = PartialCTreeRoot $ toCTreeRule <$> rules
       -- We don't validate numerical items yet
       CTree.Postlude PTAny
     toCTreeT2 T2Any = CTree.Postlude PTAny
+    toCTreeT2 (XXType2 (CTreeXXType2 v)) = absurd v
 
     toCTreeDataItem 20 =
       CTree.Literal $ Value (VBool False) mempty
@@ -224,32 +228,32 @@ buildRefCTree rules = PartialCTreeRoot $ toCTreeRule <$> rules
     toCTreeDataItem _ =
       CTree.Postlude PTAny
 
-    toCTreeGroupEntry :: GroupEntry -> CTree OrReferenced
-    toCTreeGroupEntry (GroupEntry (Just occi) _ (GEType mmkey t0)) =
+    toCTreeGroupEntry :: GroupEntry CTreePhase -> CTree OrReferenced
+    toCTreeGroupEntry (GroupEntry (Just occi) (GEType mmkey t0) _) =
       CTree.Occur
         { CTree.item = toKVPair mmkey t0
         , CTree.occurs = occi
         }
-    toCTreeGroupEntry (GroupEntry Nothing _ (GEType mmkey t0)) = toKVPair mmkey t0
-    toCTreeGroupEntry (GroupEntry (Just occi) _ (GERef n margs)) =
+    toCTreeGroupEntry (GroupEntry Nothing (GEType mmkey t0) _) = toKVPair mmkey t0
+    toCTreeGroupEntry (GroupEntry (Just occi) (GERef n margs) _) =
       CTree.Occur
         { CTree.item = CTreeE $ Ref n (fromGenArgs margs)
         , CTree.occurs = occi
         }
-    toCTreeGroupEntry (GroupEntry Nothing _ (GERef n margs)) = CTreeE $ Ref n (fromGenArgs margs)
-    toCTreeGroupEntry (GroupEntry (Just occi) _ (GEGroup g)) =
+    toCTreeGroupEntry (GroupEntry Nothing (GERef n margs) _) = CTreeE $ Ref n (fromGenArgs margs)
+    toCTreeGroupEntry (GroupEntry (Just occi) (GEGroup g) _) =
       CTree.Occur
         { CTree.item = groupToGroup g
         , CTree.occurs = occi
         }
-    toCTreeGroupEntry (GroupEntry Nothing _ (GEGroup g)) = groupToGroup g
+    toCTreeGroupEntry (GroupEntry Nothing (GEGroup g) _) = groupToGroup g
 
-    fromGenArgs :: Maybe GenericArg -> [CTree OrReferenced]
+    fromGenArgs :: Maybe (GenericArg CTreePhase) -> [CTree OrReferenced]
     fromGenArgs = maybe [] (\(GenericArg xs) -> NE.toList $ fmap toCTreeT1 xs)
 
     -- Interpret a group as an enumeration. Note that we float out the
     -- choice options
-    toCTreeEnum :: Group -> CTree OrReferenced
+    toCTreeEnum :: Group CTreePhase -> CTree OrReferenced
     toCTreeEnum (CDDL.Group (a NE.:| [])) =
       CTree.Enum . CTree.Group $ toCTreeGroupEntry <$> gcGroupEntries a
     toCTreeEnum (CDDL.Group xs) =
@@ -258,13 +262,13 @@ buildRefCTree rules = PartialCTreeRoot $ toCTreeRule <$> rules
         groupEntries = fmap gcGroupEntries xs
 
     -- Embed a group in another group, again floating out the choice options
-    groupToGroup :: Group -> CTree OrReferenced
+    groupToGroup :: Group CTreePhase -> CTree OrReferenced
     groupToGroup (CDDL.Group (a NE.:| [])) =
       CTree.Group $ fmap toCTreeGroupEntry (gcGroupEntries a)
     groupToGroup (CDDL.Group xs) =
       CTree.Choice $ fmap (CTree.Group . fmap toCTreeGroupEntry) (gcGroupEntries <$> xs)
 
-    toKVPair :: Maybe MemberKey -> Type0 -> CTree OrReferenced
+    toKVPair :: Maybe (MemberKey CTreePhase) -> Type0 CTreePhase -> CTree OrReferenced
     toKVPair Nothing t0 = toCTreeT0 t0
     toKVPair (Just mkey) t0 =
       CTree.KV
@@ -275,7 +279,7 @@ buildRefCTree rules = PartialCTreeRoot $ toCTreeRule <$> rules
         }
 
     -- Interpret a group as a map. Note that we float out the choice options
-    toCTreeMap :: Group -> CTree OrReferenced
+    toCTreeMap :: Group CTreePhase -> CTree OrReferenced
     toCTreeMap (CDDL.Group (a NE.:| [])) = CTree.Map $ fmap toCTreeGroupEntry (gcGroupEntries a)
     toCTreeMap (CDDL.Group xs) =
       CTree.Choice $
@@ -283,14 +287,14 @@ buildRefCTree rules = PartialCTreeRoot $ toCTreeRule <$> rules
 
     -- Interpret a group as an array. Note that we float out the choice
     -- options
-    toCTreeArray :: Group -> CTree OrReferenced
+    toCTreeArray :: Group CTreePhase -> CTree OrReferenced
     toCTreeArray (CDDL.Group (a NE.:| [])) =
       CTree.Array $ fmap toCTreeGroupEntry (gcGroupEntries a)
     toCTreeArray (CDDL.Group xs) =
       CTree.Choice $
         fmap (CTree.Array . fmap toCTreeGroupEntry) (gcGroupEntries <$> xs)
 
-    toCTreeMemberKey :: MemberKey -> CTree OrReferenced
+    toCTreeMemberKey :: MemberKey CTreePhase -> CTree OrReferenced
     toCTreeMemberKey (MKValue v) = CTree.Literal v
     toCTreeMemberKey (MKBareword (Name n _)) = CTree.Literal (Value (VText n) mempty)
     toCTreeMemberKey (MKType t1) = toCTreeT1 t1
@@ -300,14 +304,14 @@ buildRefCTree rules = PartialCTreeRoot $ toCTreeRule <$> rules
 --------------------------------------------------------------------------------
 
 data NameResolutionFailure
-  = UnboundReference Name
-  | MismatchingArgs Name [Name]
+  = UnboundReference (Name CTreePhase)
+  | MismatchingArgs (Name CTreePhase) [Name CTreePhase]
   | ArgsToPostlude PTerm [CTree OrReferenced]
   deriving (Show)
 
-deriving instance Eq NameResolutionFailure
+deriving instance Eq (CTree.Node OrReferenced) => Eq NameResolutionFailure
 
-postludeBinding :: Map.Map Name PTerm
+postludeBinding :: Map.Map (Name CTreePhase) PTerm
 postludeBinding =
   Map.fromList
     [ (Name "bool" mempty, PTBool)
@@ -327,9 +331,9 @@ postludeBinding =
     ]
 
 data BindingEnv i j = BindingEnv
-  { global :: Map.Map Name (ProvidedParameters (CTree i))
+  { global :: Map.Map (Name CTreePhase) (ProvidedParameters (CTree i))
   -- ^ Global name bindings via 'RuleDef'
-  , local :: Map.Map Name (CTree j)
+  , local :: Map.Map (Name CTreePhase) (CTree j)
   -- ^ Local bindings for generic parameters
   }
   deriving (Generic)
@@ -340,9 +344,9 @@ type instance CTreeExt DistReferenced = DistRef
 
 data DistRef
   = -- | Reference to a generic parameter
-    GenericRef Name
+    GenericRef (Name CTreePhase)
   | -- | Reference to a rule definition, possibly with generic arguments
-    RuleRef Name [CTree DistReferenced]
+    RuleRef (Name CTreePhase) [CTree DistReferenced]
   deriving (Eq, Generic, Show)
 
 instance Hashable DistRef
@@ -403,7 +407,7 @@ data MonoReferenced
 type instance CTreeExt MonoReferenced = MonoRef (CTree MonoReferenced)
 
 newtype MonoRef a
-  = MRuleRef Name
+  = MRuleRef (Name CTreePhase)
   deriving (Functor, Show)
 
 deriving instance Show (CTree MonoReferenced)
@@ -413,7 +417,7 @@ deriving instance Show (PartialCTreeRoot MonoReferenced)
 type MonoEnv = BindingEnv DistReferenced MonoReferenced
 
 -- | We introduce additional bindings in the state
-type MonoState = Map.Map Name (CTree MonoReferenced)
+type MonoState = Map.Map (Name CTreePhase) (CTree MonoReferenced)
 
 -- | Monad to run the monomorphisation process. We need some additional
 -- capabilities for this, so 'Either' doesn't fully cut it anymore.
@@ -435,10 +439,10 @@ newtype MonoM a = MonoM
   deriving
     ( HasSource
         "local"
-        (Map.Map Name (CTree MonoReferenced))
+        (Map.Map (Name CTreePhase) (CTree MonoReferenced))
     , HasReader
         "local"
-        (Map.Map Name (CTree MonoReferenced))
+        (Map.Map (Name CTreePhase) (CTree MonoReferenced))
     )
     via Field
           "local"
@@ -452,10 +456,10 @@ newtype MonoM a = MonoM
   deriving
     ( HasSource
         "global"
-        (Map.Map Name (ProvidedParameters (CTree DistReferenced)))
+        (Map.Map (Name CTreePhase) (ProvidedParameters (CTree DistReferenced)))
     , HasReader
         "global"
-        (Map.Map Name (ProvidedParameters (CTree DistReferenced)))
+        (Map.Map (Name CTreePhase) (ProvidedParameters (CTree DistReferenced)))
     )
     via Field
           "global"
@@ -481,7 +485,7 @@ throwNR :: NameResolutionFailure -> MonoM a
 throwNR = throw @"nameResolution"
 
 -- | Synthesize a monomorphic rule definition, returning the name
-synthMono :: Name -> [CTree DistReferenced] -> MonoM Name
+synthMono :: Name CTreePhase -> [CTree DistReferenced] -> MonoM (Name CTreePhase)
 synthMono n@(Name origName _) args =
   let fresh =
         -- % is not a valid CBOR name, so this should avoid conflict
@@ -548,7 +552,7 @@ buildMonoCTree (PartialCTreeRoot ct) = do
 -- Combined resolution
 --------------------------------------------------------------------------------
 
-fullResolveCDDL :: CDDL -> Either NameResolutionFailure (CTreeRoot MonoReferenced)
+fullResolveCDDL :: CDDL CTreePhase -> Either NameResolutionFailure (CTreeRoot MonoReferenced)
 fullResolveCDDL cddl = do
   let refCTree = buildRefCTree (asMap cddl)
   rCTree <- buildResolvedCTree refCTree
diff --git a/src/Codec/CBOR/Cuddle/Comments.hs b/src/Codec/CBOR/Cuddle/Comments.hs
index 74c789f..e65135e 100644
--- a/src/Codec/CBOR/Cuddle/Comments.hs
+++ b/src/Codec/CBOR/Cuddle/Comments.hs
@@ -26,6 +26,7 @@ import Data.List.NonEmpty (NonEmpty (..))
 import Data.String (IsString (..))
 import Data.Text qualified as T
 import Data.TreeDiff (ToExpr)
+import Data.Void (Void, absurd)
 import Data.Word (Word16, Word32, Word64, Word8)
 import GHC.Generics (Generic (..), K1 (..), M1 (..), U1 (..), V1, (:*:) (..), (:+:) (..))
 import Optics.Core (Lens', lens, view, (%~), (&), (.~), (^.))
@@ -85,6 +86,10 @@ class CollectComments a where
   default collectComments :: (Generic a, GCollectComments (Rep a)) => a -> [Comment]
   collectComments = collectCommentsG . from
 
+instance CollectComments Void where collectComments = absurd
+
+instance CollectComments () where collectComments = mempty
+
 instance CollectComments a => CollectComments (Maybe a) where
   collectComments Nothing = []
   collectComments (Just x) = collectComments x
diff --git a/src/Codec/CBOR/Cuddle/Huddle.hs b/src/Codec/CBOR/Cuddle/Huddle.hs
index eb4669b..25702e3 100644
--- a/src/Codec/CBOR/Cuddle/Huddle.hs
+++ b/src/Codec/CBOR/Cuddle/Huddle.hs
@@ -21,6 +21,13 @@ module Codec.CBOR.Cuddle.Huddle (
   IsType0 (..),
   Value (..),
 
+  -- * AST extensions
+  HuddleStage,
+  C.XCddl (..),
+  C.XTerm (..),
+  C.XXTopLevel (..),
+  C.XXType2 (..),
+
   -- * Rules and assignment
   (=:=),
   (=:~),
@@ -109,6 +116,20 @@ import GHC.Generics (Generic)
 import Optics.Core (lens, view, (%~), (&), (.~), (^.))
 import Prelude hiding ((/))
 
+data HuddleStage
+
+newtype instance C.XTerm HuddleStage = HuddleXTerm C.Comment
+  deriving (Generic, Semigroup, Monoid, Show, Eq)
+
+newtype instance C.XCddl HuddleStage = HuddleXCddl [C.Comment]
+  deriving (Generic, Semigroup, Monoid, Show, Eq)
+
+newtype instance C.XXTopLevel HuddleStage = HuddleXXTopLevel C.Comment
+  deriving (Generic, Semigroup, Monoid, Show, Eq)
+
+newtype instance C.XXType2 HuddleStage = HuddleXXType2 Void
+  deriving (Generic, Semigroup, Show, Eq)
+
 data Named a = Named
   { name :: T.Text
   , value :: a
@@ -432,7 +453,7 @@ unconstrained v = Constrained (CValue v) def []
 -- | A constraint on a 'Value' is something applied via CtlOp or RangeOp on a
 -- Type2, forming a Type1.
 data ValueConstraint a = ValueConstraint
-  { applyConstraint :: C.Type2 -> C.Type1
+  { applyConstraint :: C.Type2 HuddleStage -> C.Type1 HuddleStage
   , showConstraint :: String
   }
 
@@ -462,7 +483,7 @@ instance IsSizeable CGRefType
 
 -- | Things which can be used on the RHS of the '.size' operator.
 class IsSize a where
-  sizeAsCDDL :: a -> C.Type2
+  sizeAsCDDL :: a -> C.Type2 HuddleStage
   sizeAsString :: a -> String
 
 instance IsSize Word where
@@ -1062,15 +1083,15 @@ collectFromInit rules =
 --------------------------------------------------------------------------------
 
 -- | Convert from Huddle to CDDL, generating a top level root element.
-toCDDL :: Huddle -> CDDL
+toCDDL :: Huddle -> CDDL HuddleStage
 toCDDL = toCDDL' True
 
 -- | Convert from Huddle to CDDL, skipping a root element.
-toCDDLNoRoot :: Huddle -> CDDL
+toCDDLNoRoot :: Huddle -> CDDL HuddleStage
 toCDDLNoRoot = toCDDL' False
 
 -- | Convert from Huddle to CDDL for the purpose of pretty-printing.
-toCDDL' :: Bool -> Huddle -> CDDL
+toCDDL' :: Bool -> Huddle -> CDDL HuddleStage
 toCDDL' mkPseudoRoot hdl =
   C.fromRules
     $ ( if mkPseudoRoot
@@ -1082,14 +1103,14 @@ toCDDL' mkPseudoRoot hdl =
     toCDDLItem (HIRule r) = toCDDLRule r
     toCDDLItem (HIGroup g) = toCDDLGroup g
     toCDDLItem (HIGRule g) = toGenRuleDef g
-    toTopLevelPseudoRoot :: [Rule] -> C.Rule
+    toTopLevelPseudoRoot :: [Rule] -> C.Rule HuddleStage
     toTopLevelPseudoRoot topRs =
       toCDDLRule $
         comment "Pseudo-rule introduced by Cuddle to collect root elements" $
           "huddle_root_defs" =:= arr (fromList (fmap a topRs))
-    toCDDLRule :: Rule -> C.Rule
+    toCDDLRule :: Rule -> C.Rule HuddleStage
     toCDDLRule (Named n t0 c) =
-      (\x -> C.Rule (C.Name n mempty) Nothing C.AssignEq x (foldMap C.Comment c))
+      (\x -> C.Rule (C.Name n mempty) Nothing C.AssignEq x (foldMap (HuddleXTerm . C.Comment) c))
         . C.TOGType
         . C.Type0
         $ toCDDLType1 <$> choiceToNE t0
@@ -1103,18 +1124,18 @@ toCDDL' mkPseudoRoot hdl =
     toCDDLValue' (LText t) = C.VText t
     toCDDLValue' (LBytes b) = C.VBytes b
 
-    mapToCDDLGroup :: Map -> C.Group
+    mapToCDDLGroup :: Map -> C.Group HuddleStage
     mapToCDDLGroup xs = C.Group $ mapChoiceToCDDL <$> choiceToNE xs
 
-    mapChoiceToCDDL :: MapChoice -> C.GrpChoice
+    mapChoiceToCDDL :: MapChoice -> C.GrpChoice HuddleStage
     mapChoiceToCDDL (MapChoice entries) = C.GrpChoice (fmap mapEntryToCDDL entries) mempty
 
-    mapEntryToCDDL :: MapEntry -> C.GroupEntry
+    mapEntryToCDDL :: MapEntry -> C.GroupEntry HuddleStage
     mapEntryToCDDL (MapEntry k v occ cmnt) =
       C.GroupEntry
         (toOccurrenceIndicator occ)
-        cmnt
         (C.GEType (Just $ toMemberKey k) (toCDDLType0 v))
+        (HuddleXTerm cmnt)
 
     toOccurrenceIndicator :: Occurs -> Maybe C.OccurrenceIndicator
     toOccurrenceIndicator (Occurs Nothing Nothing) = Nothing
@@ -1123,7 +1144,7 @@ toCDDL' mkPseudoRoot hdl =
     toOccurrenceIndicator (Occurs (Just 1) Nothing) = Just C.OIOneOrMore
     toOccurrenceIndicator (Occurs lb ub) = Just $ C.OIBounded lb ub
 
-    toCDDLType1 :: Type2 -> C.Type1
+    toCDDLType1 :: Type2 -> C.Type1 HuddleStage
     toCDDLType1 = \case
       T2Constrained (Constrained x constr _) ->
         -- TODO Need to handle choices at the top level
@@ -1142,28 +1163,28 @@ toCDDL' mkPseudoRoot hdl =
       T2Generic g -> C.Type1 (toGenericCall g) Nothing mempty
       T2GenericRef (GRef n) -> C.Type1 (C.T2Name (C.Name n mempty) Nothing) Nothing mempty
 
-    toMemberKey :: Key -> C.MemberKey
+    toMemberKey :: Key -> C.MemberKey HuddleStage
     toMemberKey (LiteralKey (Literal (LText t) _)) = C.MKBareword (C.Name t mempty)
     toMemberKey (LiteralKey v) = C.MKValue $ toCDDLValue v
     toMemberKey (TypeKey t) = C.MKType (toCDDLType1 t)
 
-    toCDDLType0 :: Type0 -> C.Type0
+    toCDDLType0 :: Type0 -> C.Type0 HuddleStage
     toCDDLType0 = C.Type0 . fmap toCDDLType1 . choiceToNE
 
-    arrayToCDDLGroup :: Array -> C.Group
+    arrayToCDDLGroup :: Array -> C.Group HuddleStage
     arrayToCDDLGroup xs = C.Group $ arrayChoiceToCDDL <$> choiceToNE xs
 
-    arrayChoiceToCDDL :: ArrayChoice -> C.GrpChoice
-    arrayChoiceToCDDL (ArrayChoice entries cmt) = C.GrpChoice (fmap arrayEntryToCDDL entries) cmt
+    arrayChoiceToCDDL :: ArrayChoice -> C.GrpChoice HuddleStage
+    arrayChoiceToCDDL (ArrayChoice entries cmt) = C.GrpChoice (fmap arrayEntryToCDDL entries) (HuddleXTerm cmt)
 
-    arrayEntryToCDDL :: ArrayEntry -> C.GroupEntry
+    arrayEntryToCDDL :: ArrayEntry -> C.GroupEntry HuddleStage
     arrayEntryToCDDL (ArrayEntry k v occ cmnt) =
       C.GroupEntry
         (toOccurrenceIndicator occ)
-        cmnt
         (C.GEType (fmap toMemberKey k) (toCDDLType0 v))
+        (HuddleXTerm cmnt)
 
-    toCDDLPostlude :: Value a -> C.Name
+    toCDDLPostlude :: Value a -> C.Name HuddleStage
     toCDDLPostlude VBool = C.Name "bool" mempty
     toCDDLPostlude VUInt = C.Name "uint" mempty
     toCDDLPostlude VNInt = C.Name "nint" mempty
@@ -1181,7 +1202,7 @@ toCDDL' mkPseudoRoot hdl =
       CRef r -> C.Name (name r) mempty
       CGRef (GRef n) -> C.Name n mempty
 
-    toCDDLRanged :: Ranged -> C.Type1
+    toCDDLRanged :: Ranged -> C.Type1 HuddleStage
     toCDDLRanged (Unranged x) =
       C.Type1 (C.T2Value $ toCDDLValue x) Nothing mempty
     toCDDLRanged (Ranged lb ub rop) =
@@ -1190,18 +1211,18 @@ toCDDL' mkPseudoRoot hdl =
         (Just (C.RangeOp rop, toCDDLRangeBound ub))
         mempty
 
-    toCDDLRangeBound :: RangeBound -> C.Type2
+    toCDDLRangeBound :: RangeBound -> C.Type2 HuddleStage
     toCDDLRangeBound (RangeBoundLiteral l) = C.T2Value $ toCDDLValue l
     toCDDLRangeBound (RangeBoundRef (Named n _ _)) = C.T2Name (C.Name n mempty) Nothing
 
-    toCDDLGroup :: Named Group -> C.Rule
+    toCDDLGroup :: Named Group -> C.Rule HuddleStage
     toCDDLGroup (Named n (Group t0s) c) =
       C.Rule
         (C.Name n mempty)
         Nothing
         C.AssignEq
         ( C.TOGGroup
-            . C.GroupEntry Nothing mempty
+            . (\x -> C.GroupEntry Nothing x mempty)
             . C.GEGroup
             . C.Group
             . (NE.:| [])
@@ -1210,15 +1231,15 @@ toCDDL' mkPseudoRoot hdl =
               arrayEntryToCDDL
               t0s
         )
-        (foldMap C.Comment c)
+        (foldMap (HuddleXTerm . C.Comment) c)
 
-    toGenericCall :: GRuleCall -> C.Type2
+    toGenericCall :: GRuleCall -> C.Type2 HuddleStage
     toGenericCall (Named n gr _) =
       C.T2Name
         (C.Name n mempty)
         (Just . C.GenericArg $ fmap toCDDLType1 (args gr))
 
-    toGenRuleDef :: GRuleDef -> C.Rule
+    toGenRuleDef :: GRuleDef -> C.Rule HuddleStage
     toGenRuleDef (Named n gr c) =
       C.Rule
         (C.Name n mempty)
@@ -1228,7 +1249,7 @@ toCDDL' mkPseudoRoot hdl =
             . C.Type0
             $ toCDDLType1 <$> choiceToNE (body gr)
         )
-        (foldMap C.Comment c)
+        (foldMap (HuddleXTerm . C.Comment) c)
       where
         gps =
           C.GenericParam $ fmap (\(GRef t) -> C.Name t mempty) (args gr)
diff --git a/src/Codec/CBOR/Cuddle/IndexMappable.hs b/src/Codec/CBOR/Cuddle/IndexMappable.hs
new file mode 100644
index 0000000..eba4d78
--- /dev/null
+++ b/src/Codec/CBOR/Cuddle/IndexMappable.hs
@@ -0,0 +1,225 @@
+{-# LANGUAGE DefaultSignatures #-}
+
+module Codec.CBOR.Cuddle.IndexMappable (IndexMappable (..)) where
+
+import Codec.CBOR.Cuddle.CDDL (
+  CDDL (..),
+  GenericArg (..),
+  GenericParam (..),
+  Group (..),
+  GroupEntry (..),
+  GroupEntryVariant (..),
+  GrpChoice (..),
+  MemberKey (..),
+  Name (..),
+  Rule (..),
+  TopLevel (..),
+  Type0 (..),
+  Type1 (..),
+  Type2 (..),
+  TypeOrGroup (..),
+  XCddl,
+  XTerm,
+  XXTopLevel,
+  XXType2,
+ )
+import Codec.CBOR.Cuddle.CDDL.CTree (
+  CTreePhase,
+  XCddl (..),
+  XTerm (..),
+  XXTopLevel (..),
+  XXType2 (..),
+ )
+import Codec.CBOR.Cuddle.Huddle (HuddleStage, XCddl (..), XTerm (..), XXTopLevel (..), XXType2 (..))
+import Codec.CBOR.Cuddle.Parser (ParserStage, XCddl (..), XTerm (..), XXTopLevel (..), XXType2 (..))
+import Codec.CBOR.Cuddle.Pretty (PrettyStage, XCddl (..), XTerm (..), XXTopLevel (..))
+import Data.Bifunctor (Bifunctor (..))
+import Data.Coerce (Coercible, coerce)
+import Data.Void (absurd)
+
+class IndexMappable f i j where
+  mapIndex :: f i -> f j
+  default mapIndex :: Coercible (f i) (f j) => f i -> f j
+  mapIndex = coerce
+
+instance
+  ( IndexMappable XCddl i j
+  , IndexMappable XXTopLevel i j
+  , IndexMappable XXType2 i j
+  , IndexMappable XTerm i j
+  ) =>
+  IndexMappable CDDL i j
+  where
+  mapIndex (CDDL r tls e) = CDDL (mapIndex r) (mapIndex <$> tls) (mapIndex <$> e)
+
+instance
+  ( IndexMappable XXType2 i j
+  , IndexMappable XTerm i j
+  ) =>
+  IndexMappable Rule i j
+  where
+  mapIndex (Rule n mg a t c) = Rule (mapIndex n) (mapIndex <$> mg) a (mapIndex t) (mapIndex c)
+
+instance
+  ( IndexMappable XXTopLevel i j
+  , IndexMappable XXType2 i j
+  , IndexMappable XTerm i j
+  ) =>
+  IndexMappable TopLevel i j
+  where
+  mapIndex (TopLevelRule r) = TopLevelRule $ mapIndex r
+  mapIndex (XXTopLevel e) = XXTopLevel $ mapIndex e
+
+instance IndexMappable XTerm i j => IndexMappable Name i j where
+  mapIndex (Name n e) = Name n $ mapIndex e
+
+instance IndexMappable XTerm i j => IndexMappable GenericParam i j where
+  mapIndex (GenericParam ns) = GenericParam $ mapIndex <$> ns
+
+instance
+  ( IndexMappable XXType2 i j
+  , IndexMappable XTerm i j
+  ) =>
+  IndexMappable TypeOrGroup i j
+  where
+  mapIndex (TOGType t) = TOGType $ mapIndex t
+  mapIndex (TOGGroup g) = TOGGroup $ mapIndex g
+
+instance
+  ( IndexMappable XTerm i j
+  , IndexMappable XXType2 i j
+  ) =>
+  IndexMappable GroupEntry i j
+  where
+  mapIndex (GroupEntry mo gev e) = GroupEntry mo (mapIndex gev) (mapIndex e)
+
+instance
+  ( IndexMappable XXType2 i j
+  , IndexMappable XTerm i j
+  ) =>
+  IndexMappable GroupEntryVariant i j
+  where
+  mapIndex (GEType mk t) = GEType (mapIndex <$> mk) $ mapIndex t
+  mapIndex (GERef n ma) = GERef (mapIndex n) (mapIndex <$> ma)
+  mapIndex (GEGroup g) = GEGroup (mapIndex g)
+
+instance
+  ( IndexMappable XXType2 i j
+  , IndexMappable XTerm i j
+  ) =>
+  IndexMappable MemberKey i j
+  where
+  mapIndex (MKType t) = MKType $ mapIndex t
+  mapIndex (MKBareword n) = MKBareword $ mapIndex n
+  mapIndex (MKValue x) = MKValue x
+
+instance
+  ( IndexMappable XXType2 i j
+  , IndexMappable XTerm i j
+  ) =>
+  IndexMappable Type0 i j
+  where
+  mapIndex (Type0 ts) = Type0 $ mapIndex <$> ts
+
+instance
+  ( IndexMappable XXType2 i j
+  , IndexMappable XTerm i j
+  ) =>
+  IndexMappable Type1 i j
+  where
+  mapIndex (Type1 t mo e) = Type1 (mapIndex t) (second mapIndex <$> mo) (mapIndex e)
+
+instance
+  ( IndexMappable XXType2 i j
+  , IndexMappable XTerm i j
+  ) =>
+  IndexMappable Type2 i j
+  where
+  mapIndex (T2Value v) = T2Value v
+  mapIndex (T2Name n mg) = T2Name (mapIndex n) (mapIndex <$> mg)
+  mapIndex (T2Group t) = T2Group $ mapIndex t
+  mapIndex (T2Map g) = T2Map $ mapIndex g
+  mapIndex (T2Array a) = T2Array $ mapIndex a
+  mapIndex (T2Unwrapped n mg) = T2Unwrapped (mapIndex n) (mapIndex <$> mg)
+  mapIndex (T2Enum g) = T2Enum $ mapIndex g
+  mapIndex (T2EnumRef n mg) = T2EnumRef (mapIndex n) (mapIndex <$> mg)
+  mapIndex (T2Tag mt t) = T2Tag mt $ mapIndex t
+  mapIndex (T2DataItem t mt) = T2DataItem t mt
+  mapIndex T2Any = T2Any
+  mapIndex (XXType2 e) = XXType2 $ mapIndex e
+
+instance
+  ( IndexMappable XXType2 i j
+  , IndexMappable XTerm i j
+  ) =>
+  IndexMappable GenericArg i j
+  where
+  mapIndex (GenericArg g) = GenericArg $ mapIndex <$> g
+
+instance
+  ( IndexMappable XTerm i j
+  , IndexMappable XXType2 i j
+  ) =>
+  IndexMappable Group i j
+  where
+  mapIndex (Group g) = Group $ mapIndex <$> g
+
+instance
+  ( IndexMappable XTerm i j
+  , IndexMappable XXType2 i j
+  ) =>
+  IndexMappable GrpChoice i j
+  where
+  mapIndex (GrpChoice gs e) = GrpChoice (mapIndex <$> gs) $ mapIndex e
+
+-- ParserStage -> PrettyStage
+
+instance IndexMappable XCddl ParserStage PrettyStage where
+  mapIndex (ParserXCddl c) = PrettyXCddl c
+
+instance IndexMappable XTerm ParserStage PrettyStage where
+  mapIndex (ParserXTerm c) = PrettyXTerm c
+
+instance IndexMappable XXType2 ParserStage PrettyStage where
+  mapIndex (ParserXXType2 v) = absurd v
+
+instance IndexMappable XXTopLevel ParserStage PrettyStage where
+  mapIndex (ParserXXTopLevel c) = PrettyXXTopLevel c
+
+-- ParserStage -> CTreePhase
+
+instance IndexMappable XCddl ParserStage CTreePhase where
+  mapIndex (ParserXCddl c) = CTreeXCddl c
+
+instance IndexMappable XXTopLevel ParserStage CTreePhase where
+  mapIndex (ParserXXTopLevel c) = CTreeXXTopLevel c
+
+instance IndexMappable XXType2 ParserStage CTreePhase where
+  mapIndex (ParserXXType2 c) = CTreeXXType2 c
+
+instance IndexMappable XTerm ParserStage CTreePhase where
+  mapIndex (ParserXTerm c) = CTreeXTerm c
+
+-- ParserStage -> HuddleStage
+
+instance IndexMappable XCddl ParserStage HuddleStage where
+  mapIndex (ParserXCddl c) = HuddleXCddl c
+
+instance IndexMappable XXTopLevel ParserStage HuddleStage where
+  mapIndex (ParserXXTopLevel c) = HuddleXXTopLevel c
+
+instance IndexMappable XXType2 ParserStage HuddleStage where
+  mapIndex (ParserXXType2 c) = HuddleXXType2 c
+
+instance IndexMappable XTerm ParserStage HuddleStage where
+  mapIndex (ParserXTerm c) = HuddleXTerm c
+
+-- ParserStage -> ParserStage
+
+instance IndexMappable XCddl ParserStage ParserStage
+
+instance IndexMappable XXTopLevel ParserStage ParserStage
+
+instance IndexMappable XXType2 ParserStage ParserStage
+
+instance IndexMappable XTerm ParserStage ParserStage
diff --git a/src/Codec/CBOR/Cuddle/Parser.hs b/src/Codec/CBOR/Cuddle/Parser.hs
index f21f3f3..d71c18f 100644
--- a/src/Codec/CBOR/Cuddle/Parser.hs
+++ b/src/Codec/CBOR/Cuddle/Parser.hs
@@ -1,12 +1,22 @@
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedLabels #-}
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TypeFamilies #-}
 
 module Codec.CBOR.Cuddle.Parser where
 
 import Codec.CBOR.Cuddle.CDDL
 import Codec.CBOR.Cuddle.CDDL.CtlOp (CtlOp)
 import Codec.CBOR.Cuddle.CDDL.CtlOp qualified as COp
-import Codec.CBOR.Cuddle.Comments (Comment, WithComment (..), withComment, (!*>), (//-), (<*!))
+import Codec.CBOR.Cuddle.Comments (
+  Comment,
+  HasComment (..),
+  WithComment (..),
+  withComment,
+  (!*>),
+  (//-),
+  (<*!),
+ )
 import Codec.CBOR.Cuddle.Parser.Lexer (
   Parser,
   charInRange,
@@ -16,36 +26,58 @@ import Codec.CBOR.Cuddle.Parser.Lexer (
 import Control.Applicative.Combinators.NonEmpty qualified as NE
 import Data.Foldable (Foldable (..))
 import Data.Functor (void, ($>))
-import Data.List.NonEmpty (NonEmpty)
+import Data.List.NonEmpty (NonEmpty (..))
 import Data.List.NonEmpty qualified as NE
 import Data.Maybe (isJust)
 import Data.Text (Text)
 import Data.Text qualified as T
 import Data.Text.Encoding (encodeUtf8)
+import Data.TreeDiff (ToExpr)
+import Data.Void (Void)
+import GHC.Generics (Generic)
 import GHC.Word (Word64, Word8)
+import Optics.Core ((&), (.~))
 import Text.Megaparsec
 import Text.Megaparsec.Char hiding (space)
 import Text.Megaparsec.Char qualified as C
 import Text.Megaparsec.Char.Lexer qualified as L
 
-pCDDL :: Parser CDDL
+data ParserStage
+
+newtype instance XXTopLevel ParserStage = ParserXXTopLevel Comment
+  deriving (Generic, Show, Eq, ToExpr)
+
+newtype instance XXType2 ParserStage = ParserXXType2 Void
+  deriving (Generic, Show, Eq, ToExpr)
+
+newtype instance XTerm ParserStage = ParserXTerm {unParserXTerm :: Comment}
+  deriving (Generic, Semigroup, Monoid, Show, Eq, ToExpr)
+
+newtype instance XCddl ParserStage = ParserXCddl [Comment]
+  deriving (Generic, Semigroup, Monoid, Show, Eq, ToExpr)
+
+instance HasComment (XTerm ParserStage) where
+  commentL = #unParserXTerm
+
+pCDDL :: Parser (CDDL ParserStage)
 pCDDL = do
   initialComments <- many (try $ C.space *> pCommentBlock <* notFollowedBy pRule)
   initialRuleComment <- C.space *> optional pCommentBlock
   initialRule <- pRule
   cddlTail <- many $ pTopLevel <* C.space
-  eof $> CDDL initialComments (initialRule //- fold initialRuleComment) cddlTail
+  eof
+    $> CDDL (initialRule //- fold initialRuleComment) cddlTail (ParserXXTopLevel <$> initialComments)
 
-pTopLevel :: Parser TopLevel
+pTopLevel :: Parser (TopLevel ParserStage)
 pTopLevel = try tlRule <|> tlComment
   where
     tlRule = do
       mCmt <- optional pCommentBlock
       rule <- pRule
       pure . TopLevelRule $ rule //- fold mCmt
-    tlComment = TopLevelComment <$> pCommentBlock
+    tlComment = XXTopLevel . ParserXXTopLevel <$> pCommentBlock
 
-pRule :: Parser Rule
+pRule :: Parser (Rule ParserStage)
 pRule = do
   name <- pName
   genericParam <- optcomp pGenericParam
@@ -59,9 +91,9 @@ pRule = do
             <*> (TOGType <$> pType0 <* notFollowedBy (space >> (":" <|> "=>")))
       , (,) <$> pAssignG <* space <*> (TOGGroup <$> pGrpEntry)
       ]
-  pure $ Rule name genericParam assign typeOrGrp cmt
+  pure $ Rule name genericParam assign typeOrGrp (ParserXTerm cmt)
 
-pName :: Parser Name
+pName :: Parser (Name ParserStage)
 pName = label "name" $ do
   fc <- firstChar
   rest <- many midChar
@@ -89,20 +121,20 @@ pAssignG =
     , AssignExt <$ "//="
     ]
 
-pGenericParam :: Parser GenericParam
+pGenericParam :: Parser (GenericParam ParserStage)
 pGenericParam =
   GenericParam
     <$> between "<" ">" (NE.sepBy1 (space !*> pName <*! space) ",")
 
-pGenericArg :: Parser GenericArg
+pGenericArg :: Parser (GenericArg ParserStage)
 pGenericArg =
   GenericArg
     <$> between "<" ">" (NE.sepBy1 (space !*> pType1 <*! space) ",")
 
-pType0 :: Parser Type0
+pType0 :: Parser (Type0 ParserStage)
 pType0 = Type0 <$> sepBy1' (space !*> pType1 <*! space) (try "/")
 
-pType1 :: Parser Type1
+pType1 :: Parser (Type1 ParserStage)
 pType1 = do
   v <- pType2
   rest <- optional $ do
@@ -115,15 +147,15 @@ pType1 = do
     pure (cmtFst, tyOp, cmtSnd, w)
   case rest of
     Just (cmtFst, tyOp, cmtSnd, w) ->
-      pure $ Type1 v (Just (tyOp, w)) $ cmtFst <> cmtSnd
+      pure $ Type1 v (Just (tyOp, w)) . ParserXTerm $ cmtFst <> cmtSnd
     Nothing -> pure $ Type1 v Nothing mempty
 
-pType2 :: Parser Type2
+pType2 :: Parser (Type2 ParserStage)
 pType2 =
   choice
     [ T2Value <$> pValue
     , T2Name <$> pName <*> optional pGenericArg
-    , T2Group <$> label "group" ("(" *> space !*> pType0 <*! space <* ")")
+    , T2Group <$> label "group" ("(" *> pType0Cmt <* ")")
     , T2Map <$> label "map" ("{" *> pGroup <* "}")
     , T2Array <$> label "array" ("[" *> space !*> pGroup <*! space <* "]")
     , T2Unwrapped <$> ("~" *> space !*> pName) <*> optional pGenericArg
@@ -141,11 +173,17 @@ pType2 =
             mminor <- optional ("." *> L.decimal)
             let
               pTag
-                | major == 6 = T2Tag mminor <$> ("(" *> space !*> pType0 <*! space <* ")")
+                | major == 6 = T2Tag mminor <$> ("(" *> pType0Cmt <* ")")
                 | otherwise = empty
             pTag <|> pure (T2DataItem major mminor)
           Nothing -> pure T2Any
     ]
+  where
+    pType0Cmt = do
+      pre <- space
+      Type0 (t :| ts) <- pType0
+      post <- space
+      pure . Type0 $ (t & commentL .~ (pre <> post)) :| ts
 
 pHeadNumber :: Parser Word64
 pHeadNumber = L.decimal
@@ -176,13 +214,13 @@ pCtlOp =
                 ]
         )
 
-pGroup :: Parser Group
+pGroup :: Parser (Group ParserStage)
 pGroup = Group <$> NE.sepBy1 (space !*> pGrpChoice) "//"
 
-pGrpChoice :: Parser GrpChoice
+pGrpChoice :: Parser (GrpChoice ParserStage)
 pGrpChoice = GrpChoice <$> many (space !*> pGrpEntry <*! pOptCom) <*> mempty
 
-pGrpEntry :: Parser GroupEntry
+pGrpEntry :: Parser (GroupEntry ParserStage)
 pGrpEntry = do
   occur <- optcomp pOccur
   cmt <- space
@@ -195,9 +233,9 @@ pGrpEntry = do
       , try $ withComment <$> (GERef <$> pName <*> optional pGenericArg)
       , withComment . GEGroup <$> ("(" *> space !*> pGroup <*! space <* ")")
       ]
-  pure $ GroupEntry occur (cmt <> cmt') variant
+  pure $ GroupEntry occur variant (ParserXTerm $ cmt <> cmt')
 
-pMemberKey :: Parser (WithComment MemberKey)
+pMemberKey :: Parser (WithComment (MemberKey ParserStage))
 pMemberKey =
   choice
     [ try $ do
diff --git a/src/Codec/CBOR/Cuddle/Pretty.hs b/src/Codec/CBOR/Cuddle/Pretty.hs
index 0e8b1ba..5914fbf 100644
--- a/src/Codec/CBOR/Cuddle/Pretty.hs
+++ b/src/Codec/CBOR/Cuddle/Pretty.hs
@@ -1,5 +1,7 @@
 {-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE OverloadedLabels #-}
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 
@@ -7,7 +9,7 @@ module Codec.CBOR.Cuddle.Pretty where
 
 import Codec.CBOR.Cuddle.CDDL
 import Codec.CBOR.Cuddle.CDDL.CtlOp (CtlOp)
-import Codec.CBOR.Cuddle.Comments (CollectComments (..), Comment (..), unComment)
+import Codec.CBOR.Cuddle.Comments (CollectComments (..), Comment (..), HasComment (..), unComment)
 import Codec.CBOR.Cuddle.Pretty.Columnar (
   Cell (..),
   CellAlign (..),
@@ -24,19 +26,40 @@ import Codec.CBOR.Cuddle.Pretty.Utils (renderedLen, softspace)
 import Data.ByteString.Char8 qualified as BS
 import Data.Foldable (Foldable (..))
 import Data.List.NonEmpty qualified as NE
-import Data.String (fromString)
+import Data.String (IsString, fromString)
 import Data.Text qualified as T
+import Data.TreeDiff (ToExpr)
+import Data.Void (Void, absurd)
+import GHC.Generics (Generic)
+import Optics.Core ((^.))
 import Prettyprinter
 
-instance Pretty CDDL where
+data PrettyStage
+
+newtype instance XXTopLevel PrettyStage = PrettyXXTopLevel Comment
+  deriving (Generic, CollectComments, ToExpr, Show, Eq)
+
+newtype instance XXType2 PrettyStage = PrettyXXType2 Void
+  deriving (Generic, CollectComments, ToExpr, Show, Eq)
+
+newtype instance XTerm PrettyStage = PrettyXTerm {unPrettyXTerm :: Comment}
+  deriving (Generic, CollectComments, Semigroup, Monoid, IsString, ToExpr, Show, Eq)
+
+newtype instance XCddl PrettyStage = PrettyXCddl [Comment]
+  deriving (Generic, CollectComments, ToExpr, Show, Eq)
+
+instance HasComment (XTerm PrettyStage) where
+  commentL = #unPrettyXTerm
+
+instance Pretty (CDDL PrettyStage) where
   pretty = vsep . fmap pretty . NE.toList . cddlTopLevel
 
-instance Pretty TopLevel where
-  pretty (TopLevelComment cmt) = pretty cmt
+instance Pretty (TopLevel PrettyStage) where
+  pretty (XXTopLevel (PrettyXXTopLevel cmt)) = pretty cmt
   pretty (TopLevelRule x) = pretty x <> hardline
 
-instance Pretty Name where
-  pretty (Name name cmt) = pretty name <> prettyCommentNoBreakWS cmt
+instance Pretty (Name PrettyStage) where
+  pretty (Name name (PrettyXTerm cmt)) = pretty name <> prettyCommentNoBreakWS cmt
 
 data CommentRender
   = PreComment
@@ -54,12 +77,12 @@ instance Pretty Comment where
   pretty (Comment "") = mempty
   pretty c = prettyCommentNoBreak c <> hardline
 
-type0Def :: Type0 -> Doc ann
+type0Def :: Type0 PrettyStage -> Doc ann
 type0Def t = nest 2 $ line' <> pretty t
 
-instance Pretty Rule where
+instance Pretty (Rule PrettyStage) where
   pretty (Rule n mgen assign tog cmt) =
-    pretty cmt
+    pretty (cmt ^. commentL)
       <> groupIfNoComments
         tog
         ( pretty n <> pretty mgen <+> case tog of
@@ -74,21 +97,21 @@ instance Pretty Rule where
         AssignEq -> "="
         AssignExt -> "//="
 
-instance Pretty GenericArg where
+instance Pretty (GenericArg PrettyStage) where
   pretty (GenericArg (NE.toList -> l))
     | null (collectComments l) = group . cEncloseSep "<" ">" "," $ fmap pretty l
     | otherwise = columnarListing "<" ">" "," . Columnar $ singletonRow . pretty <$> l
 
-instance Pretty GenericParam where
+instance Pretty (GenericParam PrettyStage) where
   pretty (GenericParam (NE.toList -> l))
     | null (collectComments l) = group . cEncloseSep "<" ">" "," $ fmap pretty l
     | otherwise = columnarListing "<" ">" "," . Columnar $ singletonRow . pretty <$> l
 
-instance Pretty Type0 where
+instance Pretty (Type0 PrettyStage) where
   pretty t0@(Type0 (NE.toList -> l)) =
     groupIfNoComments t0 $ columnarSepBy "/" . Columnar $ type1ToRow <$> l
     where
-      type1ToRow (Type1 t2 tyOp cmt) =
+      type1ToRow (Type1 t2 tyOp (PrettyXTerm cmt)) =
         let
           valCell = case tyOp of
             Nothing -> cellL t2
@@ -104,15 +127,15 @@ instance Pretty TyOp where
   pretty (RangeOp Closed) = ".."
   pretty (CtrlOp n) = "." <> pretty n
 
-instance Pretty Type1 where
-  pretty (Type1 t2 Nothing cmt) = groupIfNoComments t2 (pretty t2) <> prettyCommentNoBreakWS cmt
-  pretty (Type1 t2 (Just (tyop, t2')) cmt) =
+instance Pretty (Type1 PrettyStage) where
+  pretty (Type1 t2 Nothing (PrettyXTerm cmt)) = groupIfNoComments t2 (pretty t2) <> prettyCommentNoBreakWS cmt
+  pretty (Type1 t2 (Just (tyop, t2')) (PrettyXTerm cmt)) =
     groupIfNoComments t2 (pretty t2)
       <+> pretty tyop
       <+> groupIfNoComments t2' (pretty t2')
       <> prettyCommentNoBreakWS cmt
 
-instance Pretty Type2 where
+instance Pretty (Type2 PrettyStage) where
   pretty (T2Value v) = pretty v
   pretty (T2Name n mg) = pretty n <> pretty mg
   pretty (T2Group g) = cEncloseSep "(" ")" mempty [pretty g]
@@ -131,6 +154,7 @@ instance Pretty Type2 where
       Nothing -> mempty
       Just minor -> "." <> pretty minor
   pretty T2Any = "#"
+  pretty (XXType2 (PrettyXXType2 v)) = absurd v
 
 instance Pretty OccurrenceIndicator where
   pretty OIOptional = "?"
@@ -144,7 +168,7 @@ data GroupRender
   | AsArray
   | AsGroup
 
-memberKeySep :: MemberKey -> Doc ann
+memberKeySep :: MemberKey i -> Doc ann
 memberKeySep MKType {} = " => "
 memberKeySep _ = " : "
 
@@ -165,10 +189,10 @@ groupIfNoComments x
   | not (any (mempty /=) $ collectComments x) = group
   | otherwise = id
 
-columnarGroupChoice :: GrpChoice -> Columnar ann
+columnarGroupChoice :: GrpChoice PrettyStage -> Columnar ann
 columnarGroupChoice (GrpChoice ges _cmt) = Columnar grpEntryRows
   where
-    groupEntryRow (GroupEntry oi cmt gev) =
+    groupEntryRow (GroupEntry oi gev (PrettyXTerm cmt)) =
       Row $
         [maybe emptyCell (\x -> Cell (pretty x <> space) LeftAlign) oi]
           <> groupEntryVariantCells gev
@@ -179,7 +203,7 @@ columnarGroupChoice (GrpChoice ges _cmt) = Columnar grpEntryRows
     groupEntryVariantCells (GEGroup g) = [Cell (prettyGroup AsGroup g) LeftAlign, emptyCell]
     grpEntryRows = groupEntryRow <$> ges
 
-prettyGroup :: GroupRender -> Group -> Doc ann
+prettyGroup :: GroupRender -> Group PrettyStage -> Doc ann
 prettyGroup gr g@(Group (toList -> xs)) =
   groupIfNoComments g . columnarListing (lEnc <> softspace) rEnc "// " . Columnar $
     (\x -> singletonRow . groupIfNoComments x . columnarSepBy "," $ columnarGroupChoice x) <$> xs
@@ -189,10 +213,10 @@ prettyGroup gr g@(Group (toList -> xs)) =
       AsArray -> ("[", "]")
       AsGroup -> ("(", ")")
 
-instance Pretty GroupEntry where
+instance Pretty (GroupEntry PrettyStage) where
   pretty ge = prettyColumnar . columnarGroupChoice $ GrpChoice [ge] mempty
 
-instance Pretty MemberKey where
+instance Pretty (MemberKey PrettyStage) where
   pretty (MKType t1) = pretty t1
   pretty (MKBareword n) = pretty n
   pretty (MKValue v) = pretty v
diff --git a/test/Test/Codec/CBOR/Cuddle/CDDL/Examples.hs b/test/Test/Codec/CBOR/Cuddle/CDDL/Examples.hs
index 6845565..0c1fd9a 100644
--- a/test/Test/Codec/CBOR/Cuddle/CDDL/Examples.hs
+++ b/test/Test/Codec/CBOR/Cuddle/CDDL/Examples.hs
@@ -3,14 +3,14 @@
 module Test.Codec.CBOR.Cuddle.CDDL.Examples (spec) where
 
 import Codec.CBOR.Cuddle.CDDL (Value (..), ValueVariant (..))
-import Codec.CBOR.Cuddle.CDDL.CTree (CTree (..), CTreeRoot)
-import Codec.CBOR.Cuddle.CDDL.Postlude (PTerm (..))
-import Codec.CBOR.Cuddle.CDDL.Prelude (prependPrelude)
+import Codec.CBOR.Cuddle.CDDL.CTree (CTree (..), CTreeRoot, PTerm (..))
+import Codec.CBOR.Cuddle.CDDL.Postlude (appendPostlude)
 import Codec.CBOR.Cuddle.CDDL.Resolve (
   MonoReferenced,
   NameResolutionFailure (..),
   fullResolveCDDL,
  )
+import Codec.CBOR.Cuddle.IndexMappable (IndexMappable (..))
 import Codec.CBOR.Cuddle.Parser (pCDDL)
 import Data.Text.IO qualified as T
 import Test.HUnit (assertFailure)
@@ -22,9 +22,9 @@ tryValidateFile :: FilePath -> IO (Either NameResolutionFailure (CTreeRoot MonoR
 tryValidateFile filePath = do
   contents <- T.readFile filePath
   cddl <- case parse pCDDL "" contents of
-    Right x -> pure $ prependPrelude x
+    Right x -> pure $ appendPostlude x
     Left x -> fail $ "Failed to parse the file:\n" <> errorBundlePretty x
-  pure $ fullResolveCDDL cddl
+  pure . fullResolveCDDL $ mapIndex cddl
 
 validateExpectSuccess :: FilePath -> Spec
 validateExpectSuccess filePath = it ("Successfully validates " <> filePath) $ do
diff --git a/test/Test/Codec/CBOR/Cuddle/CDDL/Gen.hs b/test/Test/Codec/CBOR/Cuddle/CDDL/Gen.hs
index d907b3f..dd5ae35 100644
--- a/test/Test/Codec/CBOR/Cuddle/CDDL/Gen.hs
+++ b/test/Test/Codec/CBOR/Cuddle/CDDL/Gen.hs
@@ -1,3 +1,5 @@
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE UndecidableInstances #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# OPTIONS_GHC -Wno-orphans #-}
 
@@ -7,6 +9,8 @@ module Test.Codec.CBOR.Cuddle.CDDL.Gen () where
 import Codec.CBOR.Cuddle.CDDL
 import Codec.CBOR.Cuddle.CDDL.CtlOp
 import Codec.CBOR.Cuddle.Comments (Comment (..))
+import Codec.CBOR.Cuddle.Parser (ParserStage, XTerm (..))
+import Codec.CBOR.Cuddle.Pretty (PrettyStage, XTerm (..), XXTopLevel (..))
 import Data.ByteString (ByteString)
 import Data.ByteString qualified as BS
 import Data.List.NonEmpty qualified as NE
@@ -15,18 +19,24 @@ import Data.Text qualified as T
 import Test.QuickCheck
 import Test.QuickCheck qualified as Gen
 
-instance Arbitrary CDDL where
+instance Arbitrary (CDDL PrettyStage) where
   arbitrary = CDDL <$> arbitrary <*> arbitrary <*> arbitrary
   shrink = genericShrink
 
-instance Arbitrary TopLevel where
+deriving newtype instance Arbitrary (XXTopLevel PrettyStage)
+
+deriving newtype instance Arbitrary (XTerm PrettyStage)
+
+instance Arbitrary (TopLevel PrettyStage) where
   arbitrary =
     Gen.oneof
-      [ TopLevelComment <$> arbitrary
+      [ XXTopLevel <$> arbitrary
       , TopLevelRule <$> arbitrary
       ]
   shrink = genericShrink
 
+deriving newtype instance Arbitrary (XTerm ParserStage)
+
 instance Arbitrary T.Text where
   arbitrary = T.pack <$> arbitrary
   shrink = fmap T.pack . shrink . T.unpack
@@ -48,7 +58,7 @@ nameMidChars = nameFstChars <> ['1' .. '9'] <> ['-', '.']
 nameEndChars :: [Char]
 nameEndChars = nameFstChars <> ['1' .. '9']
 
-instance Arbitrary Name where
+instance (Monoid (XTerm i), Arbitrary (XTerm i)) => Arbitrary (Name i) where
   arbitrary =
     let veryShortListOf = resize 3 . listOf
      in do
@@ -73,15 +83,15 @@ instance Arbitrary Assign where
   arbitrary = Gen.elements [AssignEq, AssignExt]
   shrink = genericShrink
 
-instance Arbitrary GenericParam where
+instance (Monoid (XTerm i), Arbitrary (XTerm i)) => Arbitrary (GenericParam i) where
   arbitrary = GenericParam <$> nonEmpty arbitrary
   shrink (GenericParam neName) = GenericParam <$> shrinkNE neName
 
-instance Arbitrary GenericArg where
+instance (Arbitrary (XTerm i), Monoid (XTerm i)) => Arbitrary (GenericArg i) where
   arbitrary = GenericArg <$> nonEmpty arbitrary
   shrink (GenericArg neArg) = GenericArg <$> shrinkNE neArg
 
-instance Arbitrary Rule where
+instance (Monoid (XTerm i), Arbitrary (XTerm i)) => Arbitrary (Rule i) where
   arbitrary =
     Rule
       <$> arbitrary
@@ -103,7 +113,12 @@ instance Arbitrary TyOp where
       ]
   shrink = genericShrink
 
-instance Arbitrary TypeOrGroup where
+instance
+  ( Arbitrary (XTerm i)
+  , Monoid (XTerm i)
+  ) =>
+  Arbitrary (TypeOrGroup i)
+  where
   arbitrary =
     Gen.oneof
       [ TOGGroup <$> arbitrary
@@ -111,15 +126,15 @@ instance Arbitrary TypeOrGroup where
       ]
   shrink = genericShrink
 
-instance Arbitrary Type0 where
+instance (Arbitrary (XTerm i), Monoid (XTerm i)) => Arbitrary (Type0 i) where
   arbitrary = Type0 <$> nonEmpty arbitrary
   shrink (Type0 neType1) = Type0 <$> shrinkNE neType1
 
-instance Arbitrary Type1 where
+instance (Arbitrary (XTerm i), Monoid (XTerm i)) => Arbitrary (Type1 i) where
   arbitrary = Type1 <$> arbitrary <*> arbitrary <*> arbitrary
   shrink = genericShrink
 
-instance Arbitrary Type2 where
+instance (Monoid (XTerm i), Arbitrary (XTerm i)) => Arbitrary (Type2 i) where
   arbitrary =
     recursive
       Gen.oneof
@@ -138,7 +153,6 @@ instance Arbitrary Type2 where
       [ T2Group <$> arbitrary
       , T2Tag <$> arbitrary <*> arbitrary
       ]
-  shrink = genericShrink
 
 instance Arbitrary OccurrenceIndicator where
   arbitrary =
@@ -153,15 +167,15 @@ instance Arbitrary OccurrenceIndicator where
 
   shrink = genericShrink
 
-instance Arbitrary Group where
+instance (Monoid (XTerm i), Arbitrary (XTerm i)) => Arbitrary (Group i) where
   arbitrary = Group <$> nonEmpty arbitrary
   shrink (Group gr) = Group <$> shrinkNE gr
 
-instance Arbitrary GrpChoice where
+instance (Monoid (XTerm i), Arbitrary (XTerm i)) => Arbitrary (GrpChoice i) where
   arbitrary = GrpChoice <$> listOf' arbitrary <*> pure mempty
   shrink = genericShrink
 
-instance Arbitrary GroupEntryVariant where
+instance (Monoid (XTerm i), Arbitrary (XTerm i)) => Arbitrary (GroupEntryVariant i) where
   arbitrary =
     recursive
       Gen.oneof
@@ -176,15 +190,20 @@ instance Arbitrary GroupEntryVariant where
       ]
   shrink = genericShrink
 
-instance Arbitrary GroupEntry where
+instance
+  ( Arbitrary (XTerm i)
+  , Monoid (XTerm i)
+  ) =>
+  Arbitrary (GroupEntry i)
+  where
   arbitrary =
     GroupEntry
       <$> arbitrary
-      <*> pure mempty
       <*> arbitrary
+      <*> pure mempty
   shrink = genericShrink
 
-instance Arbitrary MemberKey where
+instance (Monoid (XTerm i), Arbitrary (XTerm i)) => Arbitrary (MemberKey i) where
   arbitrary =
     recursive
       Gen.oneof
diff --git a/test/Test/Codec/CBOR/Cuddle/CDDL/Parser.hs b/test/Test/Codec/CBOR/Cuddle/CDDL/Parser.hs
index 53d777e..efedf4b 100644
--- a/test/Test/Codec/CBOR/Cuddle/CDDL/Parser.hs
+++ b/test/Test/Codec/CBOR/Cuddle/CDDL/Parser.hs
@@ -4,11 +4,10 @@ module Test.Codec.CBOR.Cuddle.CDDL.Parser where
 
 import Codec.CBOR.Cuddle.CDDL
 import Codec.CBOR.Cuddle.CDDL.CtlOp qualified as CtlOp
-import Codec.CBOR.Cuddle.Comments (Comment (..))
+import Codec.CBOR.Cuddle.IndexMappable (IndexMappable (..))
 import Codec.CBOR.Cuddle.Parser
 import Codec.CBOR.Cuddle.Parser.Lexer (Parser)
-import Codec.CBOR.Cuddle.Pretty ()
-import Data.Default.Class (Default (..))
+import Codec.CBOR.Cuddle.Pretty (PrettyStage)
 import Data.List.NonEmpty (NonEmpty (..))
 import Data.Text qualified as T
 import Data.TreeDiff (ToExpr (..), ansiWlBgEditExprCompact, exprDiff)
@@ -36,11 +35,11 @@ parserSpec = do
 
 roundtripSpec :: Spec
 roundtripSpec = describe "Roundtripping should be id" $ do
-  it "Trip Name" $ trip pName
+  it "Trip Name" $ tripIndexed pName
   xit "Trip Value" $ trip pValue
-  xit "Trip Type0" $ trip pType0
-  xit "Trip GroupEntry" $ trip pGrpEntry
-  xit "Trip Rule" $ trip pRule
+  xit "Trip Type0" $ tripIndexed pType0
+  xit "Trip GroupEntry" $ tripIndexed pGrpEntry
+  xit "Trip Rule" $ tripIndexed pRule
   where
     -- We show that, for a printed CDDL document p, print (parse p) == p. Note
     -- that we do not show that parse (print p) is p for a given generated
@@ -60,6 +59,17 @@ roundtripSpec = describe "Roundtripping should be id" $ do
                 toExpr x `exprDiff` toExpr parsed
             )
             $ printed `shouldBe` printText parsed
+    tripIndexed ::
+      forall a.
+      ( IndexMappable a ParserStage PrettyStage
+      , Eq (a PrettyStage)
+      , ToExpr (a PrettyStage)
+      , Show (a PrettyStage)
+      , Pretty (a PrettyStage)
+      , Arbitrary (a PrettyStage)
+      ) =>
+      Parser (a ParserStage) -> Property
+    tripIndexed = trip . fmap (mapIndex @a @ParserStage @PrettyStage)
     printText :: Pretty a => a -> T.Text
     printText = renderStrict . layoutPretty defaultLayoutOptions . pretty
 
@@ -179,15 +189,15 @@ type2Spec = describe "type2" $ do
                     { gcGroupEntries =
                         [ GroupEntry
                             { geOccurrenceIndicator = Nothing
-                            , geComment = Comment mempty
+                            , geExt = mempty
                             , geVariant =
                                 GEType
                                   ( Just
                                       ( MKType
                                           ( Type1
-                                              { t1Main = T2Name (Name {name = "int", nameComment = Comment mempty}) Nothing
+                                              { t1Main = T2Name (Name {name = "int", nameExt = mempty}) Nothing
                                               , t1TyOp = Nothing
-                                              , t1Comment = Comment mempty
+                                              , t1Comment = mempty
                                               }
                                           )
                                       )
@@ -195,16 +205,16 @@ type2Spec = describe "type2" $ do
                                   ( Type0
                                       { t0Type1 =
                                           Type1
-                                            { t1Main = T2Name (Name {name = "string", nameComment = Comment mempty}) Nothing
+                                            { t1Main = T2Name (Name {name = "string", nameExt = mempty}) Nothing
                                             , t1TyOp = Nothing
-                                            , t1Comment = Comment mempty
+                                            , t1Comment = mempty
                                             }
                                             :| []
                                       }
                                   )
                             }
                         ]
-                    , gcComment = Comment mempty
+                    , gcComment = mempty
                     }
                     :| []
               }
@@ -218,15 +228,15 @@ type2Spec = describe "type2" $ do
                     { gcGroupEntries =
                         [ GroupEntry
                             { geOccurrenceIndicator = Just OIZeroOrMore
-                            , geComment = Comment mempty
+                            , geExt = mempty
                             , geVariant =
                                 GEType
                                   ( Just
                                       ( MKType
                                           ( Type1
-                                              { t1Main = T2Name (Name {name = "int", nameComment = Comment mempty}) Nothing
+                                              { t1Main = T2Name (Name {name = "int", nameExt = mempty}) Nothing
                                               , t1TyOp = Nothing
-                                              , t1Comment = Comment mempty
+                                              , t1Comment = mempty
                                               }
                                           )
                                       )
@@ -234,16 +244,16 @@ type2Spec = describe "type2" $ do
                                   ( Type0
                                       { t0Type1 =
                                           Type1
-                                            { t1Main = T2Name (Name {name = "string", nameComment = Comment mempty}) Nothing
+                                            { t1Main = T2Name (Name {name = "string", nameExt = mempty}) Nothing
                                             , t1TyOp = Nothing
-                                            , t1Comment = Comment mempty
+                                            , t1Comment = mempty
                                             }
                                             :| []
                                       }
                                   )
                             }
                         ]
-                    , gcComment = Comment mempty
+                    , gcComment = mempty
                     }
                     :| []
               }
@@ -257,18 +267,18 @@ type2Spec = describe "type2" $ do
                     { gcGroupEntries =
                         [ GroupEntry
                             { geOccurrenceIndicator = Nothing
-                            , geComment = Comment mempty
+                            , geExt = mempty
                             , geVariant =
                                 GEType
                                   ( Just
-                                      (MKType (Type1 {t1Main = T2Value (value $ VUInt 1), t1TyOp = Nothing, t1Comment = Comment mempty}))
+                                      (MKType (Type1 {t1Main = T2Value (value $ VUInt 1), t1TyOp = Nothing, t1Comment = mempty}))
                                   )
                                   ( Type0
                                       { t0Type1 =
                                           Type1
-                                            { t1Main = T2Name (Name {name = "string", nameComment = Comment mempty}) Nothing
+                                            { t1Main = T2Name (Name {name = "string", nameExt = mempty}) Nothing
                                             , t1TyOp = Nothing
-                                            , t1Comment = Comment mempty
+                                            , t1Comment = mempty
                                             }
                                             :| []
                                       }
@@ -276,18 +286,18 @@ type2Spec = describe "type2" $ do
                             }
                         , GroupEntry
                             { geOccurrenceIndicator = Nothing
-                            , geComment = Comment mempty
+                            , geExt = mempty
                             , geVariant =
                                 GEType
                                   ( Just
-                                      (MKType (Type1 {t1Main = T2Value (value $ VUInt 2), t1TyOp = Nothing, t1Comment = Comment mempty}))
+                                      (MKType (Type1 {t1Main = T2Value (value $ VUInt 2), t1TyOp = Nothing, t1Comment = mempty}))
                                   )
                                   ( Type0
                                       { t0Type1 =
                                           Type1
-                                            { t1Main = T2Name (Name {name = "int", nameComment = Comment mempty}) Nothing
+                                            { t1Main = T2Name (Name {name = "int", nameExt = mempty}) Nothing
                                             , t1TyOp = Nothing
-                                            , t1Comment = Comment mempty
+                                            , t1Comment = mempty
                                             }
                                             :| []
                                       }
@@ -295,25 +305,25 @@ type2Spec = describe "type2" $ do
                             }
                         , GroupEntry
                             { geOccurrenceIndicator = Nothing
-                            , geComment = Comment mempty
+                            , geExt = mempty
                             , geVariant =
                                 GEType
                                   ( Just
-                                      (MKType (Type1 {t1Main = T2Value (value $ VUInt 3), t1TyOp = Nothing, t1Comment = Comment mempty}))
+                                      (MKType (Type1 {t1Main = T2Value (value $ VUInt 3), t1TyOp = Nothing, t1Comment = mempty}))
                                   )
                                   ( Type0
                                       { t0Type1 =
                                           Type1
-                                            { t1Main = T2Name (Name {name = "bytes", nameComment = Comment mempty}) Nothing
+                                            { t1Main = T2Name (Name {name = "bytes", nameExt = mempty}) Nothing
                                             , t1TyOp = Nothing
-                                            , t1Comment = Comment mempty
+                                            , t1Comment = mempty
                                             }
                                             :| []
                                       }
                                   )
                             }
                         ]
-                    , gcComment = Comment mempty
+                    , gcComment = mempty
                     }
                     :| []
               }
@@ -328,45 +338,45 @@ type2Spec = describe "type2" $ do
                     { gcGroupEntries =
                         [ GroupEntry
                             { geOccurrenceIndicator = Nothing
-                            , geComment = Comment mempty
+                            , geExt = mempty
                             , geVariant =
                                 GEType
                                   Nothing
                                   ( Type0
                                       { t0Type1 =
                                           Type1
-                                            { t1Main = T2Name (Name {name = "int", nameComment = Comment mempty}) Nothing
+                                            { t1Main = T2Name (Name {name = "int", nameExt = mempty}) Nothing
                                             , t1TyOp = Nothing
-                                            , t1Comment = Comment mempty
+                                            , t1Comment = mempty
                                             }
                                             :| []
                                       }
                                   )
                             }
                         ]
-                    , gcComment = Comment mempty
+                    , gcComment = mempty
                     }
                     :| [ GrpChoice
                            { gcGroupEntries =
                                [ GroupEntry
                                    { geOccurrenceIndicator = Nothing
-                                   , geComment = Comment mempty
+                                   , geExt = mempty
                                    , geVariant =
                                        GEType
                                          Nothing
                                          ( Type0
                                              { t0Type1 =
                                                  Type1
-                                                   { t1Main = T2Name (Name {name = "string", nameComment = Comment mempty}) Nothing
+                                                   { t1Main = T2Name (Name {name = "string", nameExt = mempty}) Nothing
                                                    , t1TyOp = Nothing
-                                                   , t1Comment = Comment mempty
+                                                   , t1Comment = mempty
                                                    }
                                                    :| []
                                              }
                                          )
                                    }
                                ]
-                           , gcComment = Comment mempty
+                           , gcComment = mempty
                            }
                        ]
               }
@@ -381,35 +391,35 @@ type2Spec = describe "type2" $ do
                     { gcGroupEntries =
                         [ GroupEntry
                             { geOccurrenceIndicator = Nothing
-                            , geComment = Comment mempty
+                            , geExt = mempty
                             , geVariant =
                                 GEType
                                   Nothing
                                   ( Type0
                                       { t0Type1 =
-                                          Type1 {t1Main = T2Value (value $ VUInt 0), t1TyOp = Nothing, t1Comment = Comment mempty} :| []
+                                          Type1 {t1Main = T2Value (value $ VUInt 0), t1TyOp = Nothing, t1Comment = mempty} :| []
                                       }
                                   )
                             }
                         ]
-                    , gcComment = Comment mempty
+                    , gcComment = mempty
                     }
                     :| [ GrpChoice
                            { gcGroupEntries =
                                [ GroupEntry
                                    { geOccurrenceIndicator = Nothing
-                                   , geComment = Comment mempty
+                                   , geExt = mempty
                                    , geVariant =
                                        GEType
                                          Nothing
                                          ( Type0
                                              { t0Type1 =
-                                                 Type1 {t1Main = T2Value (value $ VUInt 1), t1TyOp = Nothing, t1Comment = Comment mempty} :| []
+                                                 Type1 {t1Main = T2Value (value $ VUInt 1), t1TyOp = Nothing, t1Comment = mempty} :| []
                                              }
                                          )
                                    }
                                ]
-                           , gcComment = Comment mempty
+                           , gcComment = mempty
                            }
                        ]
               }
@@ -423,18 +433,18 @@ type2Spec = describe "type2" $ do
                     { gcGroupEntries =
                         [ GroupEntry
                             { geOccurrenceIndicator = Nothing
-                            , geComment = Comment mempty
+                            , geExt = mempty
                             , geVariant =
                                 GEType
                                   Nothing
                                   ( Type0
                                       { t0Type1 =
-                                          Type1 {t1Main = T2Value (value $ VUInt 1), t1TyOp = Nothing, t1Comment = Comment mempty} :| []
+                                          Type1 {t1Main = T2Value (value $ VUInt 1), t1TyOp = Nothing, t1Comment = mempty} :| []
                                       }
                                   )
                             }
                         ]
-                    , gcComment = Comment mempty
+                    , gcComment = mempty
                     }
                     :| []
               }
@@ -448,35 +458,35 @@ type2Spec = describe "type2" $ do
                     { gcGroupEntries =
                         [ GroupEntry
                             { geOccurrenceIndicator = Nothing
-                            , geComment = Comment mempty
+                            , geExt = mempty
                             , geVariant =
                                 GEType
                                   Nothing
                                   ( Type0
                                       { t0Type1 =
-                                          Type1 {t1Main = T2Value (value $ VUInt 2), t1TyOp = Nothing, t1Comment = Comment mempty} :| []
+                                          Type1 {t1Main = T2Value (value $ VUInt 2), t1TyOp = Nothing, t1Comment = mempty} :| []
                                       }
                                   )
                             }
                         , GroupEntry
                             { geOccurrenceIndicator = Nothing
-                            , geComment = Comment mempty
+                            , geExt = mempty
                             , geVariant =
                                 GEType
                                   Nothing
                                   ( Type0
                                       { t0Type1 =
                                           Type1
-                                            { t1Main = T2Name (Name {name = "soon", nameComment = Comment mempty}) Nothing
+                                            { t1Main = T2Name (Name {name = "soon", nameExt = mempty}) Nothing
                                             , t1TyOp = Nothing
-                                            , t1Comment = Comment mempty
+                                            , t1Comment = mempty
                                             }
                                             :| []
                                       }
                                   )
                             }
                         ]
-                    , gcComment = Comment mempty
+                    , gcComment = mempty
                     }
                     :| []
               }
@@ -488,16 +498,16 @@ grpEntrySpec = describe "GroupEntry" $ do
     parse pGrpEntry "" "int"
       `shouldParse` GroupEntry
         { geOccurrenceIndicator = Nothing
-        , geComment = Comment mempty
+        , geExt = mempty
         , geVariant =
             GEType
               Nothing
               ( Type0
                   { t0Type1 =
                       Type1
-                        { t1Main = T2Name (Name {name = "int", nameComment = Comment mempty}) Nothing
+                        { t1Main = T2Name (Name {name = "int", nameExt = mempty}) Nothing
                         , t1TyOp = Nothing
-                        , t1Comment = Comment mempty
+                        , t1Comment = mempty
                         }
                         :| []
                   }
@@ -507,16 +517,16 @@ grpEntrySpec = describe "GroupEntry" $ do
     parse pGrpEntry "" "int // notConsideredHere"
       `shouldParse` GroupEntry
         { geOccurrenceIndicator = Nothing
-        , geComment = Comment mempty
+        , geExt = mempty
         , geVariant =
             GEType
               Nothing
               ( Type0
                   { t0Type1 =
                       Type1
-                        { t1Main = T2Name (Name {name = "int", nameComment = Comment mempty}) Nothing
+                        { t1Main = T2Name (Name {name = "int", nameExt = mempty}) Nothing
                         , t1TyOp = Nothing
-                        , t1Comment = Comment mempty
+                        , t1Comment = mempty
                         }
                         :| []
                   }
@@ -526,7 +536,7 @@ grpEntrySpec = describe "GroupEntry" $ do
     parse pGrpEntry "" "a<0 ... #6(0)>"
       `shouldParse` GroupEntry
         { geOccurrenceIndicator = Nothing
-        , geComment = Comment mempty
+        , geExt = mempty
         , geVariant =
             GEType
               Nothing
@@ -535,7 +545,7 @@ grpEntrySpec = describe "GroupEntry" $ do
                       Type1
                         { t1Main =
                             T2Name
-                              (Name {name = "a", nameComment = Comment mempty})
+                              (Name {name = "a", nameExt = mempty})
                               ( Just
                                   ( GenericArg
                                       ( Type1
@@ -547,18 +557,18 @@ grpEntrySpec = describe "GroupEntry" $ do
                                                     Nothing
                                                     ( Type0
                                                         { t0Type1 =
-                                                            Type1 {t1Main = T2Value (value $ VUInt 0), t1TyOp = Nothing, t1Comment = Comment mempty} :| []
+                                                            Type1 {t1Main = T2Value (value $ VUInt 0), t1TyOp = Nothing, t1Comment = mempty} :| []
                                                         }
                                                     )
                                                 )
-                                          , t1Comment = Comment mempty
+                                          , t1Comment = mempty
                                           }
                                           :| []
                                       )
                                   )
                               )
                         , t1TyOp = Nothing
-                        , t1Comment = Comment mempty
+                        , t1Comment = mempty
                         }
                         :| []
                   }
@@ -568,28 +578,31 @@ grpEntrySpec = describe "GroupEntry" $ do
     parse pGrpEntry "" "0* a"
       `shouldParse` GroupEntry
         (Just (OIBounded (Just 0) Nothing))
-        def
         ( GEType
             Nothing
             (Type0 (Type1 (T2Name (Name "a" mempty) Nothing) Nothing mempty :| []))
         )
+        mempty
 
 grpChoiceSpec :: SpecWith ()
 grpChoiceSpec = describe "GroupChoice" $ do
   it "Should parse part of a group alternative" $
     parse pGrpChoice "" "int // string"
       `shouldParse` GrpChoice
-        [ GroupEntry Nothing mempty $
-            GEType
-              Nothing
-              ( Type0
-                  ( Type1
-                      (T2Name (Name "int" mempty) Nothing)
-                      Nothing
-                      mempty
-                      :| []
-                  )
-              )
+        [ GroupEntry
+            Nothing
+            ( GEType
+                Nothing
+                ( Type0
+                    ( Type1
+                        (T2Name (Name "int" mempty) Nothing)
+                        Nothing
+                        mempty
+                        :| []
+                    )
+                )
+            )
+            mempty
         ]
         mempty
 
@@ -629,27 +642,27 @@ qcFoundSpec =
       Type1
         { t1Main =
             T2Map
-              (Group {unGroup = GrpChoice {gcGroupEntries = [], gcComment = Comment mempty} :| []})
+              (Group {unGroup = GrpChoice {gcGroupEntries = [], gcComment = mempty} :| []})
         , t1TyOp =
             Just
               ( CtrlOp CtlOp.Ge
               , T2EnumRef
-                  (Name {name = "i", nameComment = Comment mempty})
+                  (Name {name = "i", nameExt = mempty})
                   ( Just
                       ( GenericArg
                           ( Type1
                               { t1Main =
                                   T2Map
-                                    (Group {unGroup = GrpChoice {gcGroupEntries = [], gcComment = Comment mempty} :| []})
+                                    (Group {unGroup = GrpChoice {gcGroupEntries = [], gcComment = mempty} :| []})
                               , t1TyOp = Nothing
-                              , t1Comment = Comment mempty
+                              , t1Comment = mempty
                               }
-                              :| [Type1 {t1Main = T2Value (value $ VUInt 3), t1TyOp = Nothing, t1Comment = Comment mempty}]
+                              :| [Type1 {t1Main = T2Value (value $ VUInt 3), t1TyOp = Nothing, t1Comment = mempty}]
                           )
                       )
                   )
               )
-        , t1Comment = Comment mempty
+        , t1Comment = mempty
         }
     parseExample "S = 0* ()" pRule $
       Rule
@@ -657,8 +670,10 @@ qcFoundSpec =
         Nothing
         AssignEq
         ( TOGGroup
-            ( GroupEntry (Just (OIBounded (Just 0) Nothing)) mempty $
-                GEGroup (Group (GrpChoice mempty mempty :| []))
+            ( GroupEntry
+                (Just (OIBounded (Just 0) Nothing))
+                (GEGroup (Group (GrpChoice mempty mempty :| [])))
+                mempty
             )
         )
         mempty
@@ -672,10 +687,11 @@ qcFoundSpec =
           ( TOGGroup
               ( GroupEntry
                   Nothing
+                  ( GEType
+                      (Just (MKValue (value $ VText "6 ybe2ddl8frq0vqa8zgrk07khrljq7p plrufpd1sff3p95")))
+                      (Type0 (Type1 (T2Value (value $ VText "u")) Nothing mempty :| []))
+                  )
                   mempty
-                  $ GEType
-                    (Just (MKValue (value $ VText "6 ybe2ddl8frq0vqa8zgrk07khrljq7p plrufpd1sff3p95")))
-                    (Type0 (Type1 (T2Value (value $ VText "u")) Nothing mempty :| []))
               )
           )
           mempty
diff --git a/test/Test/Codec/CBOR/Cuddle/CDDL/Pretty.hs b/test/Test/Codec/CBOR/Cuddle/CDDL/Pretty.hs
index 3a5a354..3666418 100644
--- a/test/Test/Codec/CBOR/Cuddle/CDDL/Pretty.hs
+++ b/test/Test/Codec/CBOR/Cuddle/CDDL/Pretty.hs
@@ -22,7 +22,7 @@ import Codec.CBOR.Cuddle.CDDL (
   ValueVariant (..),
   value,
  )
-import Codec.CBOR.Cuddle.Pretty ()
+import Codec.CBOR.Cuddle.Pretty (PrettyStage)
 import Data.List.NonEmpty (NonEmpty (..))
 import Data.Text qualified as T
 import Data.TreeDiff (ToExpr (..), prettyExpr)
@@ -40,13 +40,13 @@ prettyPrintsTo x s = assertEqual (show . prettyExpr $ toExpr x) s rendered
   where
     rendered = renderString (layoutPretty defaultLayoutOptions (pretty x))
 
-t2Name :: Type2
+t2Name :: Type2 PrettyStage
 t2Name = T2Name (Name "a" mempty) mempty
 
-t1Name :: Type1
+t1Name :: Type1 PrettyStage
 t1Name = Type1 t2Name Nothing mempty
 
-mkType0 :: Type2 -> Type0
+mkType0 :: Type2 PrettyStage -> Type0 PrettyStage
 mkType0 t2 = Type0 $ Type1 t2 Nothing mempty :| []
 
 spec :: Spec
@@ -56,14 +56,14 @@ spec = describe "Pretty printer" $ do
 
 qcSpec :: Spec
 qcSpec = describe "QuickCheck" $ do
-  xprop "CDDL prettyprinter leaves no trailing spaces" $ \(cddl :: CDDL) -> do
+  xprop "CDDL prettyprinter leaves no trailing spaces" $ \(cddl :: CDDL PrettyStage) -> do
     let
       prettyStr = T.pack . renderString . layoutPretty defaultLayoutOptions $ pretty cddl
       stripLines = T.unlines . fmap T.stripEnd . T.lines
     counterexample (show . prettyExpr $ toExpr cddl) $
       prettyStr `shouldBe` stripLines prettyStr
 
-drep :: Rule
+drep :: Rule PrettyStage
 drep =
   Rule
     "drep"
@@ -77,37 +77,37 @@ drep =
                         ( GrpChoice
                             [ GroupEntry
                                 Nothing
-                                mempty
                                 (GEType Nothing (Type0 $ Type1 (T2Value . value $ VUInt 0) Nothing mempty :| []))
+                                mempty
                             , GroupEntry
                                 Nothing
-                                mempty
                                 (GEType Nothing (Type0 $ Type1 (T2Name "addr_keyhash" Nothing) Nothing mempty :| []))
+                                mempty
                             ]
                             mempty
                             :| [ GrpChoice
                                    [ GroupEntry
                                        Nothing
-                                       mempty
                                        (GEType Nothing (Type0 $ Type1 (T2Value . value $ VUInt 1) Nothing mempty :| []))
+                                       mempty
                                    , GroupEntry
                                        Nothing
-                                       mempty
                                        (GEType Nothing (Type0 $ Type1 (T2Name "script_hash" Nothing) Nothing mempty :| []))
+                                       mempty
                                    ]
                                    mempty
                                , GrpChoice
                                    [ GroupEntry
                                        Nothing
-                                       mempty
                                        (GEType Nothing (Type0 $ Type1 (T2Value . value $ VUInt 2) Nothing mempty :| []))
+                                       mempty
                                    ]
                                    mempty
                                , GrpChoice
                                    [ GroupEntry
                                        Nothing
-                                       mempty
                                        (GEType Nothing (Type0 $ Type1 (T2Value . value $ VUInt 3) Nothing mempty :| []))
+                                       mempty
                                    ]
                                    mempty
                                ]
@@ -125,22 +125,22 @@ drep =
 unitSpec :: Spec
 unitSpec = describe "HUnit" $ do
   describe "Name" $ do
-    it "names" $ Name "a" mempty `prettyPrintsTo` "a"
+    it "names" $ Name @PrettyStage "a" "" `prettyPrintsTo` "a"
   describe "Type0" $ do
-    it "name" $ Type0 (t1Name :| []) `prettyPrintsTo` "a"
+    it "name" $ Type0 @PrettyStage (t1Name :| []) `prettyPrintsTo` "a"
   describe "Type1" $ do
     it "name" $ t1Name `prettyPrintsTo` "a"
   describe "Type2" $ do
     it "T2Name" $ t2Name `prettyPrintsTo` "a"
     describe "T2Array" $ do
-      let groupEntryName = GroupEntry Nothing mempty $ GERef (Name "a" mempty) Nothing
+      let groupEntryName = GroupEntry Nothing (GERef (Name "a" mempty) Nothing) ""
       it "one element" $
         T2Array (Group (GrpChoice [groupEntryName] mempty :| [])) `prettyPrintsTo` "[a]"
       it "two elements" $
         T2Array
           ( Group
               ( GrpChoice
-                  [ GroupEntry Nothing mempty $ GEType Nothing (mkType0 . T2Value . value $ VUInt 1)
+                  [ GroupEntry Nothing (GEType Nothing (mkType0 . T2Value . value $ VUInt 1)) ""
                   , groupEntryName
                   ]
                   mempty
@@ -152,8 +152,8 @@ unitSpec = describe "HUnit" $ do
         T2Array
           ( Group
               ( GrpChoice
-                  [ GroupEntry Nothing "one" $ GEType Nothing (mkType0 . T2Value . value $ VUInt 1)
-                  , GroupEntry Nothing "two" $ GEType Nothing (mkType0 . T2Value . value $ VUInt 2)
+                  [ GroupEntry Nothing (GEType Nothing (mkType0 . T2Value . value $ VUInt 1)) "one"
+                  , GroupEntry Nothing (GEType Nothing (mkType0 . T2Value . value $ VUInt 2)) "two"
                   ]
                   mempty
                   :| []
@@ -164,9 +164,14 @@ unitSpec = describe "HUnit" $ do
         T2Array
           ( Group
               ( GrpChoice
-                  [ GroupEntry Nothing "first\nmultiline comment" $ GEType Nothing (mkType0 . T2Value . value $ VUInt 1)
-                  , GroupEntry Nothing "second\nmultiline comment" $
-                      GEType Nothing (mkType0 . T2Value . value $ VUInt 2)
+                  [ GroupEntry
+                      Nothing
+                      (GEType Nothing (mkType0 . T2Value . value $ VUInt 1))
+                      "first\nmultiline comment"
+                  , GroupEntry
+                      Nothing
+                      (GEType Nothing (mkType0 . T2Value . value $ VUInt 2))
+                      "second\nmultiline comment"
                   ]
                   mempty
                   :| []
@@ -175,7 +180,7 @@ unitSpec = describe "HUnit" $ do
           `prettyPrintsTo` "[ 1 ; first\n    ; multiline comment\n, 2 ; second\n    ; multiline comment\n]"
   describe "Rule" $ do
     it "simple assignment" $
-      Rule
+      Rule @PrettyStage
         (Name "a" mempty)
         Nothing
         AssignEq
diff --git a/test/Test/Codec/CBOR/Cuddle/Huddle.hs b/test/Test/Codec/CBOR/Cuddle/Huddle.hs
index 06a0dbf..f0d42b3 100644
--- a/test/Test/Codec/CBOR/Cuddle/Huddle.hs
+++ b/test/Test/Codec/CBOR/Cuddle/Huddle.hs
@@ -5,8 +5,9 @@
 
 module Test.Codec.CBOR.Cuddle.Huddle where
 
-import Codec.CBOR.Cuddle.CDDL (CDDL, sortCDDL)
+import Codec.CBOR.Cuddle.CDDL (CDDL, fromRules, sortCDDL)
 import Codec.CBOR.Cuddle.Huddle
+import Codec.CBOR.Cuddle.IndexMappable (IndexMappable (..))
 import Codec.CBOR.Cuddle.Parser
 import Data.Text qualified as T
 import Test.Codec.CBOR.Cuddle.CDDL.Pretty qualified as Pretty
@@ -155,10 +156,10 @@ shouldMatchParse ::
 shouldMatchParse x parseFun input = parse parseFun "" (T.pack input) `shouldParse` x
 
 shouldMatchParseCDDL ::
-  CDDL ->
+  CDDL HuddleStage ->
   String ->
   Expectation
-shouldMatchParseCDDL x = shouldMatchParse x pCDDL
+shouldMatchParseCDDL x = shouldMatchParse x . fmap mapIndex $ pCDDL
 
-toSortedCDDL :: Huddle -> CDDL
-toSortedCDDL = sortCDDL . toCDDLNoRoot
+toSortedCDDL :: Huddle -> CDDL HuddleStage
+toSortedCDDL = fromRules . sortCDDL . toCDDLNoRoot