Change 'Expression' to use numbers with units.

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

3 files changed, 87 insertions, 48 deletions

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"