Change 'Expression' to use numbers with units.

This helps to catch bugs in the compiler, specifically ones related to
mixing up bits and bytes.

12 files changed, 254 insertions, 125 deletions

diff --git a/goal.fiddle b/goal.fiddle
index 24203b5..bddd3bf 100644
--- a/goal.fiddle
+++ b/goal.fiddle
@@ -130,7 +130,7 @@ package stm32l4.gpio {
 
         assert_pos(0x12);
         wo reg alt_r2(8);
-        reserved(1);
+        reserved(8);
       };
     };
     assert_pos(0x14);
diff --git a/src/Language/Fiddle/Ast/Internal/MetaTypes.hs b/src/Language/Fiddle/Ast/Internal/MetaTypes.hs
index b2fdb26..ec5948d 100644
--- a/src/Language/Fiddle/Ast/Internal/MetaTypes.hs
+++ b/src/Language/Fiddle/Ast/Internal/MetaTypes.hs
@@ -12,13 +12,18 @@ module Language.Fiddle.Ast.Internal.MetaTypes
     Guaranteed(..),
     guarantee,
     guaranteeM,
-    revokeGuarantee
+    revokeGuarantee,
+    Variant(..),
+    foldVariant,
+    toEither
   )
 where
 
 import Data.Functor.Identity as X
 import Data.List.NonEmpty as X (NonEmpty (..))
 import Data.Maybe (fromMaybe)
+import Data.Bifunctor
+import Data.Aeson (ToJSON(..))
 
 -- | 'IsMaybe' is a typeclass for converting a general functor into a Maybe.
 class IsMaybe f where
@@ -171,3 +176,54 @@ instance IsMaybe (Guaranteed s) where
 instance IsIdentity (Guaranteed True) where
   unwrap (Guaranteed t) = t
   wrap = Guaranteed
+
+-- | 'Variant' is a type-level conditional construct that holds a value of type
+--   'ifTrue' when the type-level boolean 'b' is 'True', or a value of type
+--   'ifFalse' when 'b' is 'False'.
+--
+--   This data type allows encoding conditional behavior at the type level, 
+--   where the actual type stored depends on a compile-time boolean. It provides
+--   a way to handle cases where certain fields or structures may vary based on
+--   type-level conditions.
+--
+--   The type parameters are:
+--   * 'b' - A type-level boolean that determines which type the 'Variant' holds.
+--   * 'ifTrue' - The type of the value stored when 'b' is 'True'.
+--   * 'ifFalse' - The type of the value stored when 'b' is 'False'.
+--
+--   Constructors:
+--   * 'LeftV' - Holds a value of type 'ifTrue' when 'b' is 'True'.
+--   * 'RightV' - Holds a value of type 'ifFalse' when 'b' is 'False'.
+data Variant (b :: Bool) ifTrue ifFalse where
+  -- | Constructor for the 'Variant' when the type-level boolean is 'True'. It
+  --   stores a value of type 'ifTrue'.
+  LeftV :: ifTrue -> Variant True ifTrue ifFalse
+
+  -- | Constructor for the 'Variant' when the type-level boolean is 'False'. It
+  --   stores a value of type 'ifFalse'.
+  RightV :: ifFalse -> Variant False ifTrue ifFalse
+
+toEither :: Variant b t f -> Either t f
+toEither (LeftV l) = Left l
+toEither (RightV l) = Right l
+
+foldVariant :: (t -> r) -> (f -> r) -> Variant b t f -> r
+foldVariant fl fr = either fl fr . toEither
+
+instance Functor (Variant b t) where
+  fmap _ (LeftV x) = LeftV x
+  fmap f (RightV x) = RightV (f x)
+
+instance Foldable (Variant b t) where
+  foldMap f t = foldMap f (toEither t)
+  
+instance Bifunctor (Variant b) where
+  bimap fl _ (LeftV l) = LeftV (fl l)
+  bimap _ fr (RightV r) = RightV (fr r)
+  
+instance IsIdentity (Variant False t) where
+  unwrap (RightV r) = r
+  wrap = RightV
+
+instance (ToJSON f, ToJSON t) => ToJSON (Variant s t f) where
+  toJSON = foldVariant toJSON toJSON
diff --git a/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs b/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
index f467141..ec62c84 100644
--- a/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
+++ b/src/Language/Fiddle/Ast/Internal/SyntaxTree.hs
@@ -14,8 +14,7 @@ module Language.Fiddle.Ast.Internal.SyntaxTree
     RegisterOffset,
     BitsOffset,
     QMd,
-    N (..),
-    Unit (..),
+    NamedUnit(..),
     FieldSpan (..),
     -- Witness Types
     Witness (..),
@@ -52,7 +51,6 @@ module Language.Fiddle.Ast.Internal.SyntaxTree
     mapDirectedM,
     asDirected,
     undirected,
-    bitsToBytes,
   )
 where
 
@@ -71,29 +69,11 @@ import Language.Fiddle.Ast.Internal.Kinds
 import Language.Fiddle.Ast.Internal.MetaTypes
 import Language.Fiddle.Ast.Internal.Stage
 import Language.Fiddle.Internal.UnitInterface
+import Language.Fiddle.Internal.UnitNumbers
 
 type QMd s t = When (s .>= Qualified) t
 
--- | Phantom type used to ensure Bits and Bytes don't get mixed up in the code.
-data Unit = Bits | Bytes
-
--- | An integer with a unit.
-newtype N (u :: Unit) = N Word32
-  deriving newtype (Real, Enum, Num, Eq, Ord, Integral)
-
-instance (Show (N u)) where
-  show (N b) = show b
-
-instance (ToJSON (N u)) where
-  toJSON (N b) = toJSON b
-
-instance (FromJSON (N u)) where
-  parseJSON v = N <$> parseJSON v
-
-bitsToBytes :: N Bits -> (N Bytes, N Bits)
-bitsToBytes (N a) = let (y, i) = divMod a 8 in (N y, N i)
-
-data FieldSpan (u :: Unit) where
+data FieldSpan u where
   FieldSpan ::
     { offset :: N u,
       size :: N u
@@ -112,8 +92,8 @@ type family RegisterOffset stage where
 -- | The Type of number during each stage of compilation. When in the first stage,
 -- numbers are just strings like anything else. In later stages, numbers get
 -- parsed into actual integers. This makes it easier to process later.
-type family NumberType (a :: Stage) :: Type where
-  NumberType s = If (s .< Expanded) Text Integer
+type family NumberType (u :: k) (a :: Stage) :: Type where
+  NumberType u s = If (s .< Expanded) Text (N u)
 
 -- A Name is multiple identifiers separated by dots. It's the way of namespaces
 -- to different packages.
@@ -271,15 +251,15 @@ data Identifier f a = Identifier
   deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Expressions used within Fiddle, including literals and variables.
-data Expression (s :: Stage) :: SynTree where
+data Expression (u :: unit) (s :: Stage) :: SynTree where
   -- | A numeric literal, whose value is dependent on the compilation stage.
   LitNum ::
     { -- | The numeric value.
-      litNumValue :: NumberType stage,
+      litNumValue :: Variant (stage .< Expanded) Text (N u),
       -- | Annotation for the literal.
       litNumAnnot :: a
     } ->
-    Expression stage f a
+    Expression u stage f a
   -- | A variable reference.
   Var ::
     { -- | The identifier of the variable.
@@ -287,7 +267,7 @@ data Expression (s :: Stage) :: SynTree where
       -- | Annotation for the variable.
       varAnnot :: a
     } ->
-    Expression stage f a
+    Expression u stage f a
   deriving (Generic, Annotated, Alter, Typeable, Walk)
 
 -- | Represents an import statement in the Fiddle language.
@@ -371,7 +351,7 @@ data FiddleDecl :: StagedSynTree where
       -- | The location identifier.
       locationIdent :: Identifier f a,
       -- | The associated expression.
-      locationExpr :: Expression stage f a,
+      locationExpr :: Expression Address stage f a,
       -- | Annotation for the location declaration.
       locationAnnot :: a
     } ->
@@ -407,7 +387,7 @@ data FiddleDecl :: StagedSynTree where
       -- | The identifier of the object.
       objectIdent :: Identifier f a,
       -- | The location expression.
-      objectLocation :: Expression stage f a,
+      objectLocation :: Expression Address stage f a,
       -- | The type of the object.
       objectType :: ObjType stage f a,
       -- | Annotation for the object declaration.
@@ -448,7 +428,7 @@ data ObjType stage f a where
     { -- | The type of the array elements.
       arrayObjType :: ObjType stage f a,
       -- | The size of the array.
-      arraySize :: Expression stage f a,
+      arraySize :: Expression Unitless stage f a,
       -- | Annotation for the array type.
       arrayAnnot :: a
     } ->
@@ -472,7 +452,7 @@ data ObjTypeDecl stage f a where
     { -- | Witness for stage constraint.
       disableAssertStatementsAfterConsistencyCheck :: Witness (stage .< Checked),
       -- | The expression for the assertion.
-      assertExpr :: Expression stage f a,
+      assertExpr :: Expression Bytes stage f a,
       -- | Annotation for the assertion.
       assertAnnot :: a
     } ->
@@ -489,7 +469,7 @@ data ObjTypeDecl stage f a where
       -- doesn't exist.
       regIdent :: Guaranteed (stage .>= Qualified) (Identifier f a),
       -- | Size of the register.
-      regSize :: Expression stage f a,
+      regSize :: Expression Bits stage f a,
       -- | Optional register body.
       regBody :: Maybe (RegisterBody stage f a),
       -- | Annotation for the register declaration.
@@ -503,7 +483,7 @@ data ObjTypeDecl stage f a where
       -- | Generated identifier for this reserved field.
       reservedIdent :: When (stage .>= Qualified) String,
       -- | The expression for reserved space.
-      reservedExpr :: Expression stage f a,
+      reservedExpr :: Expression Bits stage f a,
       -- | Annotation for the reserved declaration.
       reservedAnnot :: a
     } ->
@@ -579,7 +559,7 @@ data RegisterBitsDecl stage f a where
   -- | Declaration for reserved bits.
   ReservedBits ::
     { -- | Expression for reserved bits.
-      reservedBitsExpr :: Expression stage f a,
+      reservedBitsExpr :: Expression Bits stage f a,
       -- | Annotation for the reserved bits.
       reservedBitsAnnot :: a
     } ->
@@ -620,7 +600,7 @@ data RegisterBitsTypeRef stage f a where
     { -- | Reference to the array type.
       bitsArrayTypeRef :: RegisterBitsTypeRef stage f a,
       -- | Size of the array.
-      bitsArraySize :: Expression stage f a,
+      bitsArraySize :: Expression Unitless stage f a,
       -- | Annotation for the array.
       bitsArrayAnnot :: a
     } ->
@@ -648,7 +628,7 @@ data RegisterBitsTypeRef stage f a where
   -- | A direct specification of bits as an expression.
   RegisterBitsJustBits ::
     { -- | Expression for the bits.
-      justBitsExpr :: Expression stage f a,
+      justBitsExpr :: Expression Bits stage f a,
       -- | Annotation for the bits.
       justBitsAnnot :: a
     } ->
@@ -659,7 +639,7 @@ data RegisterBitsTypeRef stage f a where
 data AnonymousBitsType stage f a where
   AnonymousEnumBody ::
     { -- | Expression defining the enum size.
-      anonEnumExpr :: Expression stage f a,
+      anonEnumSize :: Expression Bits stage f a,
       -- | The body of the enum.
       anonEnumBody :: f (EnumBody stage f a),
       -- | Annotation for the anonymous enum.
@@ -673,7 +653,7 @@ data BitType (stage :: Stage) (f :: Type -> Type) a where
   -- | An enumeration type for bits.
   EnumBitType ::
     { -- | Expression defining the enum size.
-      enumBitExpr :: Expression stage f a,
+      enumBitSize :: Expression Bits stage f a,
       -- | The body of the enum.
       enumBitBody :: f (EnumBody stage f a),
       -- | Annotation for the enumeration.
@@ -683,7 +663,7 @@ data BitType (stage :: Stage) (f :: Type -> Type) a where
   -- | A raw bit type.
   RawBits ::
     { -- | Expression defining the bits.
-      rawBitsExpr :: Expression stage f a,
+      rawBitsExpr :: Expression Bits stage f a,
       -- | Annotation for the raw bits.
       rawBitsAnnot :: a
     } ->
@@ -708,7 +688,7 @@ data EnumConstantDecl stage f a where
     { -- | Identifier for the constant.
       enumConstIdent :: Identifier f a,
       -- | Expression defining the constant.
-      enumConstExpr :: Expression stage f a,
+      enumConstExpr :: Expression Unitless stage f a,
       -- | Annotation for the constant.
       enumConstAnnot :: a
     } ->
@@ -716,7 +696,7 @@ data EnumConstantDecl stage f a where
   -- | A reserved value in the enum.
   EnumConstantReserved ::
     { -- | Expression for the reserved value.
-      enumReservedExpr :: Expression stage f a,
+      enumReservedExpr :: Expression Unitless stage f a,
       -- | Annotation for the reserved value.
       enumReservedAnnot :: a
     } ->
diff --git a/src/Language/Fiddle/Ast/Internal/Util.hs b/src/Language/Fiddle/Ast/Internal/Util.hs
index 2a03227..87fca96 100644
--- a/src/Language/Fiddle/Ast/Internal/Util.hs
+++ b/src/Language/Fiddle/Ast/Internal/Util.hs
@@ -4,8 +4,11 @@ import Data.Functor.Identity
 import Data.List.NonEmpty hiding (map)
 import qualified Data.Text as Text
 import Language.Fiddle.Ast.Internal.Instances
+import Language.Fiddle.Ast.Internal.MetaTypes
+import Language.Fiddle.Ast.Internal.Stage
 import Language.Fiddle.Ast.Internal.SyntaxTree
 import Language.Fiddle.Internal.UnitInterface
+import Language.Fiddle.Internal.UnitNumbers
 import Language.Fiddle.Types
 
 -- | 'squeeze' traverses a structure ('t') with a monadic functor ('f')
@@ -81,9 +84,9 @@ directiveToMetadata (Directed directives t a) qualifiedPath =
         elements
 
 expressionToInt ::
-  (Integral i, Integral (NumberType stage)) =>
-  Expression stage f a ->
-  Either String i
+  (stage .< Expanded ~ False) =>
+  Expression u stage f a ->
+  Either String (N u)
 expressionToInt = \case
-  LitNum {litNumValue = v} -> return (fromIntegral v)
+  LitNum {litNumValue = (RightV v)} -> return (fromIntegral v)
   _ -> Left "Incorrect Expression"
diff --git a/src/Language/Fiddle/Compiler/ConsistencyCheck.hs b/src/Language/Fiddle/Compiler/ConsistencyCheck.hs
index e0c7876..4be2912 100644
--- a/src/Language/Fiddle/Compiler/ConsistencyCheck.hs
+++ b/src/Language/Fiddle/Compiler/ConsistencyCheck.hs
@@ -16,11 +16,11 @@ import qualified Data.IntMap as IntMap
 import Data.List (intercalate)
 import qualified Data.List.NonEmpty as NonEmpty
 import Data.Typeable
-import Data.Word (Word32)
 import GHC.TypeError as TypeError
 import Language.Fiddle.Ast
 import Language.Fiddle.Compiler
 import Language.Fiddle.Internal.UnitInterface as UnitInterface
+import Language.Fiddle.Internal.UnitNumbers
 import Language.Fiddle.Types
 import Text.Printf (printf)
 import Prelude hiding (unzip)
@@ -114,7 +114,7 @@ instance AdvanceStage S FiddleUnit where
         Maybe (t S f a)
       castTS = cast
 
-deriving instance AdvanceStage S Expression
+deriving instance AdvanceStage S (Expression u)
 
 deriving instance AdvanceStage S RegisterBitsTypeRef
 
@@ -122,7 +122,7 @@ deriving instance AdvanceStage S ObjType
 
 deriving instance (AdvanceStage S t) => AdvanceStage S (Directed t)
 
-advanceObjTypeBody :: ObjTypeBody S F A -> Word32 -> M (Word32, ObjTypeBody S' F A)
+advanceObjTypeBody :: ObjTypeBody S F A -> N Bytes -> M (N Bytes, ObjTypeBody S' F A)
 advanceObjTypeBody (ObjTypeBody us decls a) startOffset = do
   (decls', _) <- advanceDecls
 
@@ -134,13 +134,19 @@ advanceObjTypeBody (ObjTypeBody us decls a) startOffset = do
 
   return (calcSize, ObjTypeBody us (reverse $ map snd decls') a)
   where
-    advanceDecls :: M ([(Word32, Directed ObjTypeDecl S' F A)], Word32)
+    advanceDecls :: M ([(N Bytes, Directed ObjTypeDecl S' F A)], N Bytes)
     advanceDecls = do
       foldlM
         ( \(ret, offset) d ->
-            let advanceOffset = case us of
+            let advanceOffset :: N Bytes -> N Bytes -> N Bytes
+                advanceOffset = case us of
                   Union {} -> const
                   Struct {} -> (+)
+                doReturn ::
+                  (Monad m) =>
+                  ObjTypeDecl S' F A ->
+                  N Bytes ->
+                  m ([(N Bytes, Directed ObjTypeDecl S' F A)], N Bytes)
                 doReturn x size = return ((size, mapDirected (const x) d) : ret, advanceOffset offset size)
              in case undirected d of
                   e@AssertPosStatement {assertExpr = expr} -> do
@@ -149,28 +155,35 @@ advanceObjTypeBody (ObjTypeBody us decls a) startOffset = do
                     return (ret, offset)
                   (RegisterDecl _ mod ident size Nothing a) -> do
                     (sizeExpr, reifiedSize) <- advanceAndGetSize size
-                    let span = Present (FieldSpan (N offset) (N reifiedSize))
+                    reifiedSizeBytes <- checkBitsSizeMod8 a reifiedSize
+
+                    let span = Present (FieldSpan offset reifiedSizeBytes)
                     doReturn (RegisterDecl span mod ident sizeExpr Nothing a)
                       =<< checkBitsSizeMod8 a reifiedSize
                   (RegisterDecl _ mod ident size (Just body) a) -> do
                     declaredSize <- expressionToIntM size
                     (actualSize, body') <- advanceRegisterBody body
+
                     checkSizeMismatch a declaredSize actualSize
+
                     (sizeExpr, reifiedSize) <- advanceAndGetSize size
-                    let span = Present (FieldSpan (N offset) (N reifiedSize))
-                    doReturn (RegisterDecl span mod ident sizeExpr (Just body') a)
-                      =<< checkBitsSizeMod8 a reifiedSize
+                    reifiedSizeBytes <- checkBitsSizeMod8 a reifiedSize
+
+                    let span = Present (FieldSpan offset reifiedSizeBytes)
+                    doReturn (RegisterDecl span mod ident sizeExpr (Just body') a) reifiedSizeBytes
                   (ReservedDecl _ i size a) -> do
                     (sizeExpr, reifiedSize) <- advanceAndGetSize size
-                    let span = Present (FieldSpan (N offset) (N reifiedSize))
-                    doReturn (ReservedDecl span i sizeExpr a) reifiedSize
+                    reifiedSizeBytes <- checkBitsSizeMod8 a reifiedSize
+                    let span = Present (FieldSpan offset reifiedSizeBytes)
+                    doReturn (ReservedDecl span i sizeExpr a) reifiedSizeBytes
                   (TypeSubStructure (Identity body) name a) -> do
                     (size, body') <- advanceObjTypeBody body offset
                     doReturn (TypeSubStructure (Identity body') name a) size
         )
-        (([], startOffset) :: ([(Word32, Directed ObjTypeDecl S' F A)], Word32))
+        (([], startOffset) :: ([(N Bytes, Directed ObjTypeDecl S' F A)], N Bytes))
         decls
 
+    advanceAndGetSize :: Expression u S F A -> M (Expression u S' F A, N u)
     advanceAndGetSize e = (,) <$> advanceStage () e <*> expressionToIntM e
 
 pattern RegisterBodyPattern :: BodyType F A -> [Directed RegisterBitsDecl s F A] -> A -> A -> RegisterBody s F A
@@ -178,7 +191,7 @@ pattern RegisterBodyPattern u decls a b = RegisterBody u (Identity (DeferredRegi
 
 -- registerBodyPattern u decls a b = RegisterBody u (Identity (DeferredRegisterBody decls a)) a
 
-advanceRegisterBody :: RegisterBody S F A -> M (Word32, RegisterBody S' F A)
+advanceRegisterBody :: RegisterBody S F A -> M (N Bits, RegisterBody S' F A)
 -- Handle the case where it's a union.
 advanceRegisterBody
   (RegisterBodyPattern us (NonEmpty.nonEmpty -> Just decls) a b) = do
@@ -207,7 +220,7 @@ advanceRegisterBody (RegisterBodyPattern u _ a b) =
   return (0, RegisterBodyPattern u [] a b)
 advanceRegisterBody RegisterBody {} = error "GHC not smart enuf"
 
-checkJagged :: (Annotated t) => [(Word32, t f A)] -> Compile s ()
+checkJagged :: (Annotated t) => [(N u, t f A)] -> Compile s ()
 checkJagged decls = do
   let expectedSize = maximum (fmap fst decls)
   forM_ decls $ \(sz, annot -> a) ->
@@ -223,7 +236,7 @@ checkJagged decls = do
         )
         a
 
-advanceDecl :: Word32 -> RegisterBitsDecl S F A -> M (Word32, RegisterBitsDecl S' F A)
+advanceDecl :: N Bits -> RegisterBitsDecl S F A -> M (N Bits, RegisterBitsDecl S' F A)
 advanceDecl offset = \case
   ReservedBits expr an -> do
     sz <- expressionToIntM expr
@@ -234,18 +247,18 @@ advanceDecl offset = \case
           )
   DefinedBits _ mod ident typ annot -> do
     size <- bitsTypeSize typ
-    let span = Present (FieldSpan (N offset) (N size))
+    let span = Present (FieldSpan offset size)
     (size,)
       <$> (DefinedBits span mod ident <$> advanceStage () typ <*> pure annot)
   BitsSubStructure subBody subName ann -> do
     (sz, body') <- advanceRegisterBody subBody
     return (sz, BitsSubStructure body' subName ann)
 
-bitsTypeSize :: RegisterBitsTypeRef S F A -> M Word32
+bitsTypeSize :: RegisterBitsTypeRef S F A -> M (N Bits)
 bitsTypeSize (RegisterBitsArray tr nExpr _) = do
   sz <- bitsTypeSize tr
   n <- expressionToIntM nExpr
-  return (sz * n)
+  return (sz .*. n)
 bitsTypeSize
   RegisterBitsReference
     { bitsRefQualificationMetadata =
@@ -255,18 +268,18 @@ bitsTypeSize (RegisterBitsReference {}) = error "should be exhaustive"
 bitsTypeSize (RegisterBitsJustBits expr _) =
   expressionToIntM expr
 
-checkSizeMismatch :: A -> Word32 -> Word32 -> Compile s ()
+checkSizeMismatch :: (NamedUnit u) => A -> N u -> N u -> Compile s ()
 checkSizeMismatch _ a b | a == b = return ()
 checkSizeMismatch pos declaredSize calculatedSize =
   emitDiagnosticError
     ( printf
-        "Size assertion failed. Declared size %d, calculated %d"
-        declaredSize
-        calculatedSize
+        "Size assertion failed. Declared size %s, calculated %s"
+        (unitName declaredSize)
+        (unitName calculatedSize)
     )
     pos
 
-checkPositionAssertion :: A -> Word32 -> Word32 -> Compile s ()
+checkPositionAssertion :: A -> N u -> N u -> Compile s ()
 checkPositionAssertion _ a b | a == b = return ()
 checkPositionAssertion pos declaredPosition calculatedPostion =
   emitDiagnosticError
@@ -278,9 +291,9 @@ checkPositionAssertion pos declaredPosition calculatedPostion =
     pos
 
 expressionToIntM ::
-  (Integral i, Integral (NumberType stage)) =>
-  Expression stage f A ->
-  Compile s i
+  (stage .< Expanded ~ False) =>
+  Expression u stage f A ->
+  Compile s (N u)
 expressionToIntM expr =
   resolveOrFail $
     either
@@ -289,21 +302,22 @@ expressionToIntM expr =
       return
       (expressionToInt expr)
 
-checkBitsSizeMod8 :: A -> Word32 -> M Word32
-checkBitsSizeMod8 _ w | w `mod` 8 == 0 = return (w `div` 8)
+checkBitsSizeMod8 :: A -> N Bits -> M (N Bytes)
 checkBitsSizeMod8 a w = do
-  emitDiagnosticWarning
-    (printf "Register size %d is not a multiple of 8. Please add padding to this register." w)
-    a
-  return ((w `div` 8) + 1)
-
-checkEnumConsistency :: Expression S F A -> EnumBody S F A -> M ()
+  let (x, rem) = bitsToBytes w
+  when (rem /= 0) $
+    emitDiagnosticError
+      (printf "Register size %d is not a multiple of 8. Please add padding to this register." w)
+      a
+  return x
+
+checkEnumConsistency :: Expression Bits S F A -> EnumBody S F A -> M ()
 checkEnumConsistency expr enumBody@(EnumBody {enumConsts = constants}) = do
   declaredSize <- expressionToIntM expr
 
   -- If the declared size is less than or equal to 4, we'll enforce that the
   -- enum is packed. This is to make sure the user has covered all bases.
-  when (declaredSize <= (4 :: Word32)) $ do
+  when (declaredSize <= 4) $ do
     imap <-
       foldlM
         ( \imap (undirected -> enumConst) -> do
@@ -311,7 +325,7 @@ checkEnumConsistency expr enumBody@(EnumBody {enumConsts = constants}) = do
               EnumConstantDecl _ expr _ -> expressionToIntM expr
               EnumConstantReserved expr _ -> expressionToIntM expr
 
-            when (number >= (2 :: Word32) ^ declaredSize) $
+            when (number >= 2 ^ declaredSize) $
               emitDiagnosticError
                 ( printf
                     "Enum constant too large. Max allowed %d\n"
diff --git a/src/Language/Fiddle/Compiler/Expansion.hs b/src/Language/Fiddle/Compiler/Expansion.hs
index 11a68be..d9bce1e 100644
--- a/src/Language/Fiddle/Compiler/Expansion.hs
+++ b/src/Language/Fiddle/Compiler/Expansion.hs
@@ -8,17 +8,17 @@ module Language.Fiddle.Compiler.Expansion (expandAst, expansionPhase) where
 
 import Control.Monad.Identity (Identity (..))
 import Control.Monad.State (get, modify, put)
+import qualified Data.Char as Char
 import Data.List (intercalate)
+import qualified Data.List.NonEmpty as NonEmpty
 import Data.Text (Text)
+import qualified Data.Text as Text
 import Language.Fiddle.Ast
 import Language.Fiddle.Compiler
 import Language.Fiddle.Compiler.Qualification ()
+import Language.Fiddle.Internal.UnitNumbers
 import Language.Fiddle.Types
 
-import qualified Data.Char as Char
-import qualified Data.List.NonEmpty as NonEmpty
-import qualified Data.Text as Text
-
 type M = Compile State
 
 type Annot = Commented SourceSpan
@@ -108,10 +108,11 @@ instance AdvanceStage CurrentStage FiddleUnit where
   advanceStage path (FiddleUnit v decls a) =
     FiddleUnit v <$> reconfigureFiddleDecls path decls <*> pure a
 
-instance AdvanceStage CurrentStage Expression where
+instance AdvanceStage CurrentStage (Expression u) where
   advanceStage _ = \case
     (Var i a) -> return $ Var i a
-    (LitNum t a) -> LitNum <$> parseNum (unCommented a) t <*> pure a
+    (LitNum (LeftV t) a) ->
+      LitNum . RightV <$> parseNum (unCommented a) t <*> pure a
 
 instance AdvanceStage CurrentStage RegisterBitsTypeRef where
   advanceStage path = \case
@@ -146,18 +147,19 @@ instance AdvanceStage CurrentStage ObjType where
         <*> advanceStage path expr
         <*> pure a
 
-parseNum :: SourceSpan -> Text -> Compile s Integer
-parseNum span txt = fromMayberOrFail span "Unable to parse number" $
-  case Text.unpack (Text.take 2 txt) of
-    "0b" -> toNumWithRadix (Text.drop 2 txt) 2
-    "0x" -> toNumWithRadix (Text.drop 2 txt) 16
-    ('0' : _) -> toNumWithRadix (Text.tail txt) 8
-    _ -> toNumWithRadix txt 10
+parseNum :: SourceSpan -> Text -> Compile s (N u)
+parseNum span txt = fmap NumberWithUnit $
+  fromMayberOrFail span "Unable to parse number" $
+    case Text.unpack (Text.take 2 txt) of
+      "0b" -> toNumWithRadix (Text.drop 2 txt) 2
+      "0x" -> toNumWithRadix (Text.drop 2 txt) 16
+      ('0' : _) -> toNumWithRadix (Text.tail txt) 8
+      _ -> toNumWithRadix txt 10
   where
     removeUnders :: Text -> Text
     removeUnders = Text.replace (Text.pack "_") Text.empty
 
-    toNumWithRadix :: Text -> Int -> Maybe Integer
+    toNumWithRadix :: Text -> Int -> Maybe Int
     toNumWithRadix (removeUnders -> txt) (fromIntegral -> radix) =
       Text.foldl
         ( \mAcc x ->
@@ -166,7 +168,7 @@ parseNum span txt = fromMayberOrFail span "Unable to parse number" $
         (Just 0)
         txt
 
-    digitToInt :: Char -> Integer -> Maybe Integer
+    digitToInt :: Char -> Int -> Maybe Int
     digitToInt (Char.toLower -> ch) radix =
       let a
             | Char.isDigit ch = Just (Char.ord ch - Char.ord '0')
diff --git a/src/Language/Fiddle/Compiler/ImportResolution.hs b/src/Language/Fiddle/Compiler/ImportResolution.hs
index 2249714..a27a1dc 100644
--- a/src/Language/Fiddle/Compiler/ImportResolution.hs
+++ b/src/Language/Fiddle/Compiler/ImportResolution.hs
@@ -123,7 +123,7 @@ deriving instance AdvanceStage CurrentStage ObjTypeDecl
 
 deriving instance AdvanceStage CurrentStage FiddleUnit
 
-deriving instance AdvanceStage CurrentStage Expression
+deriving instance AdvanceStage CurrentStage (Expression u)
 
 deriving instance AdvanceStage CurrentStage RegisterBitsTypeRef
 
diff --git a/src/Language/Fiddle/Compiler/Qualification.hs b/src/Language/Fiddle/Compiler/Qualification.hs
index 70378c3..ce6250a 100644
--- a/src/Language/Fiddle/Compiler/Qualification.hs
+++ b/src/Language/Fiddle/Compiler/Qualification.hs
@@ -24,6 +24,7 @@ import Language.Fiddle.Compiler
 import Language.Fiddle.Compiler.ConsistencyCheck ()
 import Language.Fiddle.Internal.Scopes
 import Language.Fiddle.Internal.UnitInterface as UnitInterface
+import Language.Fiddle.Internal.UnitNumbers
 import Language.Fiddle.Types
 import Text.Printf (printf)
 
@@ -113,7 +114,7 @@ deriving instance AdvanceStage S RegisterBitsDecl
 
 deriving instance AdvanceStage S ObjTypeDecl
 
-deriving instance AdvanceStage S Expression
+deriving instance AdvanceStage S (Expression u)
 
 instance AdvanceStage S RegisterBitsTypeRef where
   advanceStage localState = \case
@@ -332,7 +333,7 @@ objTypeToExport ls = \case
     (full, _ :: ExportedTypeDecl) <- resolveOrFail =<< resolveName n ls
     return $ ReferencedObjectType (intercalate "." full)
 
-calculateTypeSize :: ObjTypeBody Expanded F A -> M Word32
+calculateTypeSize :: ObjTypeBody Expanded F A -> M (N Bytes)
 calculateTypeSize (ObjTypeBody bodyType decls _) =
   ( case bodyType of
       Union {} -> maximum
@@ -340,33 +341,33 @@ calculateTypeSize (ObjTypeBody bodyType decls _) =
   )
     <$> mapM calculateDeclSize decls
   where
-    calculateDeclSize :: Directed ObjTypeDecl Expanded F A -> M Word32
+    calculateDeclSize :: Directed ObjTypeDecl Expanded F A -> M (N Bytes)
     calculateDeclSize (undirected -> decl) =
       case decl of
         AssertPosStatement {} -> return 0
-        RegisterDecl {regSize = size} -> expressionToIntM size
-        ReservedDecl {reservedExpr = size} -> expressionToIntM size
+        RegisterDecl {regSize = size} -> fst . bitsToBytes <$> expressionToIntM size
+        ReservedDecl {reservedExpr = size} -> fst . bitsToBytes <$> expressionToIntM size
         TypeSubStructure {subStructureBody = b} -> calculateTypeSize =<< resolveOrFail b
 
-getBitTypeDeclaredSize :: BitType Expanded F A -> M Word32
+getBitTypeDeclaredSize :: BitType Expanded F A -> M (N Bits)
 getBitTypeDeclaredSize = \case
   (EnumBitType declaredSize _ _) -> expressionToIntM declaredSize
   (RawBits declaredSize _) -> expressionToIntM declaredSize
 
 resolveLocationExpression ::
-  (Integral i, Integral (NumberType stage)) =>
+  (stage .< Expanded ~ False) =>
   LocalState ->
-  Expression stage F A ->
-  M i
+  Expression u stage F A ->
+  M (N u)
 resolveLocationExpression ls (Var var _) = do
   (_, ExportedLocationDecl _ v) <- resolveOrFail =<< resolveName var ls
   return (fromIntegral v)
 resolveLocationExpression _ e = expressionToIntM e
 
 expressionToIntM ::
-  (Integral i, Integral (NumberType stage)) =>
-  Expression stage f A ->
-  M i
+  (stage .< Expanded ~ False) =>
+  Expression u stage f A ->
+  M (N u)
 expressionToIntM expr =
   resolveOrFail $
     either
diff --git a/src/Language/Fiddle/GenericTree.hs b/src/Language/Fiddle/GenericTree.hs
index 9b04ac7..1dfafc2 100644
--- a/src/Language/Fiddle/GenericTree.hs
+++ b/src/Language/Fiddle/GenericTree.hs
@@ -28,9 +28,7 @@ import Language.Fiddle.Internal.UnitInterface
 import Language.Fiddle.Types
 
 type Context stage =
-  ( Show (NumberType stage),
-    Typeable stage,
-    ToJSON (NumberType stage),
+  ( Typeable stage,
     ToJSON (RegisterOffset stage)
   )
 
@@ -277,4 +275,4 @@ deriving instance (Context stage) => (ToGenericSyntaxTree (EnumConstantDecl stag
 
 deriving instance (Context stage) => (ToGenericSyntaxTree (PackageBody stage))
 
-deriving instance (Context stage) => (ToGenericSyntaxTree (Expression stage))
+deriving instance (Context stage) => (ToGenericSyntaxTree (Expression u stage))
diff --git a/src/Language/Fiddle/Internal/UnitInterface.hs b/src/Language/Fiddle/Internal/UnitInterface.hs
index c5cbc2c..2a538eb 100644
--- a/src/Language/Fiddle/Internal/UnitInterface.hs
+++ b/src/Language/Fiddle/Internal/UnitInterface.hs
@@ -8,6 +8,7 @@ import GHC.Generics
 import Language.Fiddle.Internal.Scopes (Scope)
 import qualified Language.Fiddle.Internal.Scopes as Scopes
 import Language.Fiddle.Types (SourceSpan)
+import Language.Fiddle.Internal.UnitNumbers
 
 data InternalDirectiveExpression
   = InternalDirectiveExpressionNumber String
@@ -93,7 +94,7 @@ data ExportedLocationDecl where
     { -- | Metadata associated with the location.
       exportedLocationMetadata :: Metadata,
       -- | The value of the location as an integer.
-      exportedLocationValue :: Integer
+      exportedLocationValue :: N Address
     } ->
     ExportedLocationDecl
   deriving (Show, Eq, Ord, Generic, FromJSON, ToJSON)
@@ -105,7 +106,7 @@ data ExportedBitsDecl where
     { -- | Metadata associated with the bits declaration.
       exportedBitsDeclMetadata :: Metadata,
       -- | The size of the bits in this declaration.
-      exportedBitsDeclSizeBits :: Word32
+      exportedBitsDeclSizeBits :: N Bits
     } ->
     ExportedBitsDecl
   deriving (Show, Eq, Ord, Generic, FromJSON, ToJSON)
@@ -117,7 +118,7 @@ data ExportedTypeDecl where
     { -- | Metadata associated with the type declaration.
       exportedTypeDeclMetadata :: Metadata,
       -- | The size of the type in bytes.
-      exportedTypeDeclSizeBytes :: Word32
+      exportedTypeDeclSizeBytes :: N Bytes
     } ->
     ExportedTypeDecl
   deriving (Show, Eq, Ord, Generic, FromJSON, ToJSON)
@@ -127,7 +128,7 @@ data ReferencedObjectType where
     {objectTypeReference :: String} -> ReferencedObjectType
   ArrayObjectType ::
     { arrayObjectTypeType :: ReferencedObjectType,
-      arryObjecttTypeNumber :: Word32
+      arryObjecttTypeNumber :: N Unitless
     } ->
     ReferencedObjectType
   deriving (Show, Eq, Ord, Generic, FromJSON, ToJSON)
@@ -139,7 +140,7 @@ data ExportedObjectDecl where
     { -- | Metadata associated with the object declaration.
       exportedObjectDeclMetadata :: Metadata,
       -- | The memory location of the object.
-      exportedObjectDeclLocation :: Integer,
+      exportedObjectDeclLocation :: N Address,
       -- | The type of the object as a string.
       exportedObjectDeclType :: ReferencedObjectType
     } ->
diff --git a/src/Language/Fiddle/Internal/UnitNumbers.hs b/src/Language/Fiddle/Internal/UnitNumbers.hs
new file mode 100644
index 0000000..7bdc539
--- /dev/null
+++ b/src/Language/Fiddle/Internal/UnitNumbers.hs
@@ -0,0 +1,74 @@
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+
+module Language.Fiddle.Internal.UnitNumbers where
+
+import Data.Aeson (FromJSON (..), ToJSON (..))
+import Text.Printf
+
+-- | 'NumberWithUnit' represents a numeric value associated with a unit type.
+--   The unit type 'u' is a phantom type, meaning it carries no runtime data
+--   but provides compile-time type safety for the units.
+--   The actual value is of type 'i', such as 'Int' or 'Integer'.
+newtype NumberWithUnit u i = NumberWithUnit i
+  deriving newtype (Real, Enum, Num, Eq, Ord, Integral, PrintfArg)
+
+-- | Custom 'Show' instance for 'NumberWithUnit' to display the numeric value.
+instance (Show i) => Show (NumberWithUnit u i) where
+  show (NumberWithUnit b) = show b
+
+-- | JSON serialization for 'NumberWithUnit', converting to and from JSON using
+--   the underlying numeric type.
+instance (ToJSON i) => ToJSON (NumberWithUnit u i) where
+  toJSON (NumberWithUnit b) = toJSON b
+
+instance (FromJSON i) => FromJSON (NumberWithUnit u i) where
+  parseJSON v = NumberWithUnit <$> parseJSON v
+
+-- | Type alias for 'NumberWithUnit' where the underlying numeric type is 'Int'.
+--   This is a shorthand for convenience when working with integers.
+type N u = NumberWithUnit u Int
+
+-- | Phantom types representing different units that can be used with
+--   'NumberWithUnit'. These types carry no data but serve as compile-time
+--   markers for the units.
+data Bits
+data Bytes
+data Address
+data Unitless
+
+-- | 'NamedUnit' is a typeclass for units that can be represented as strings.
+--   It provides a 'unitName' method for displaying a formatted representation
+--   of the unit.
+class NamedUnit u where
+  unitName :: N u -> String
+
+-- | 'NamedUnit' instance for 'Bits'. It displays the value followed by "bit"
+--   or "bits", depending on whether the value is 1 or not.
+instance NamedUnit Bits where
+  unitName 1 = "1 bit"
+  unitName n = show n ++ " bits"
+
+-- | 'NamedUnit' instance for 'Bytes'. It displays the value followed by "byte"
+--   or "bytes", depending on whether the value is 1 or not.
+instance NamedUnit Bytes where
+  unitName 1 = "1 byte"
+  unitName n = show n ++ " bytes"
+
+-- | Convert a value from 'Bits' to 'Bytes'. This returns a tuple containing the
+--   number of complete bytes and the remaining bits.
+bitsToBytes :: N Bits -> (N Bytes, N Bits)
+bitsToBytes (NumberWithUnit a) =
+  let (y, i) = divMod a 8 in (NumberWithUnit y, NumberWithUnit i)
+
+-- | Convert a value from 'Bytes' to 'Bits'. This multiplies the byte value by
+--   8 to get the corresponding number of bits.
+bytesToBits :: N Bytes -> N Bits
+bytesToBits (NumberWithUnit a) = NumberWithUnit (a * 8)
+
+-- | Multiplication operator for 'NumberWithUnit' values. The result retains
+--   the original unit 'u', while the right operand is of the 'Unitless' type.
+--   This operator is useful for scaling values by a unitless factor.
+(.*.) :: N u -> N Unitless -> N u
+(.*.) (NumberWithUnit a) (NumberWithUnit b) = NumberWithUnit (a * b)
+
diff --git a/src/Language/Fiddle/Parser.hs b/src/Language/Fiddle/Parser.hs
index 8d2eab2..8070c1c 100644
--- a/src/Language/Fiddle/Parser.hs
+++ b/src/Language/Fiddle/Parser.hs
@@ -190,7 +190,7 @@ objTypeP = do
                 AnonymousObjType Witness <$> defer body (objTypeBodyP t)
             )
 
-exprInParenP :: Pa Expression
+exprInParenP :: Pa (Expression u)
 exprInParenP = tok TokLParen *> expressionP <* tok TokRParen
 
 objTypeBodyP :: BodyType F (Commented SourceSpan) -> Pa ObjTypeBody
@@ -311,12 +311,12 @@ enumConstantDeclP =
     (tok KWReserved >> EnumConstantReserved <$> (tok TokEq >> expressionP))
       <|> (EnumConstantDecl <$> ident <*> (tok TokEq >> expressionP))
 
-expressionP :: Pa Expression
+expressionP :: Pa (Expression u)
 expressionP =
   withMeta $
     token
       ( \case
-          (TokLitNum num) -> Just (LitNum num)
+          (TokLitNum num) -> Just (LitNum $ LeftV num)
           _ -> Nothing
       )
       <|> (Var <$> name)