Change 'Expression' to use numbers with units.

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

4 files changed, 81 insertions, 64 deletions

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