Compile-time evaluate constant BitArray int segments on JavaScript

CHANGELOG.md

@@ -207,6 +207,11 @@
 
   ([Richard Viney](https://github.com/richard-viney))
 
+- When targeting JavaScript the compiler now generates faster and smaller code
+  for `Int` values in bit array expressions and patterns by evaluating them at
+  compile time where possible.
+  ([Richard Viney](https://github.com/richard-viney))
+
 ### Formatter
 
 - The formatter no longer removes the first argument from a function

compiler-core/Cargo.toml

@@ -46,6 +46,7 @@ unicode-segmentation = "1.12.0"
 bimap = "0.6.3"
 # Parsing of arbitrary width int values
 num-bigint = "0.4.6"
+num-traits = "0.2.19"
 async-trait.workspace = true
 base16.workspace = true
 bytes.workspace = true

compiler-core/src/javascript.rs

@@ -6,6 +6,9 @@ mod pattern;
 mod tests;
 mod typescript;
 
+use num_bigint::BigInt;
+use num_traits::{One, ToPrimitive};
+
 use crate::analyse::TargetSupport;
 use crate::build::Target;
 use crate::codegen::TypeScriptDeclarations;
@@ -785,3 +788,56 @@ fn bool(bool: bool) -> Document<'static> {
         false => "false".to_doc(),
     }
 }
+
+/// Int segments <= 48 bits wide in bit arrays are within JavaScript's safe range and are evaluated
+/// at compile time when all inputs are known. This is done for both bit array expressions and
+/// pattern matching.
+///
+/// Int segments of any size could be evaluated at compile time, but currently aren't due to the
+/// potential for causing large generated JS for inputs such as `<<0:8192>>`.
+///
+pub(crate) const SAFE_INT_SEGMENT_MAX_SIZE: usize = 48;
+
+/// Evaluates the value of an Int segment in a bit array into its corresponding bytes. This avoids
+/// needing to do the evaluation at runtime when all inputs are known at compile-time.
+///
+pub(crate) fn bit_array_segment_int_value_to_bytes(
+    mut value: BigInt,
+    size: BigInt,
+    endianness: endianness::Endianness,
+    location: SrcSpan,
+) -> Result<Vec<u8>, Error> {
+    // Clamp negative sizes to zero
+    let size = size.max(BigInt::ZERO);
+
+    // The segment size in bits is not allowed to exceed the range of a u32. At runtime the
+    // limit is lower than this and depends on the JS engine. V8's limit is currently
+    // 2^30-1 bits.
+    let size = size.to_u32().ok_or_else(|| Error::Unsupported {
+        feature: "Integer segment size greater than 2^32-1".into(),
+        location,
+    })?;
+
+    // Convert negative number to two's complement representation
+    if value < BigInt::ZERO {
+        let value_modulus = BigInt::one() << size;
+        value = &value_modulus + (value % &value_modulus);
+    }
+
+    let byte_mask = BigInt::from(0xFF);
+
+    // Convert value to the desired number of bytes
+    let mut bytes = vec![0u8; size as usize / 8];
+    for byte in bytes.iter_mut() {
+        *byte = (&value & &byte_mask)
+            .to_u8()
+            .expect("bitwise and result to be a u8");
+        value >>= 8;
+    }
+
+    if endianness.is_big() {
+        bytes.reverse();
+    }
+
+    Ok(bytes)
+}

compiler-core/src/javascript/expression.rs

@@ -1,3 +1,4 @@
+use num_bigint::BigInt;
 use vec1::Vec1;
 
 use super::{
@@ -232,19 +233,36 @@ impl<'module> Generator<'module> {
                 let details = self.sized_bit_array_segment_details(segment)?;
 
                 if segment.type_ == crate::type_::int() {
-                    if details.has_explicit_size {
-                        self.tracker.sized_integer_segment_used = true;
-                        Ok(docvec![
-                            "sizedInt(",
-                            value,
-                            ", ",
-                            details.size,
-                            ", ",
-                            bool(details.endianness.is_big()),
-                            ")"
-                        ])
-                    } else {
-                        Ok(value)
+                    match (details.size_value, segment.value.as_ref()) {
+                        (Some(size_value), TypedExpr::Int { int_value, .. })
+                            if size_value <= SAFE_INT_SEGMENT_MAX_SIZE.into() =>
+                        {
+                            let bytes = bit_array_segment_int_value_to_bytes(
+                                int_value.clone(),
+                                size_value,
+                                details.endianness,
+                                segment.location,
+                            )?;
+
+                            Ok(u8_slice(&bytes))
+                        }
+
+                        (Some(size_value), _) if size_value == 8.into() => Ok(value),
+
+                        (Some(size_value), _) if size_value <= 0.into() => Ok(docvec![]),
+
+                        _ => {
+                            self.tracker.sized_integer_segment_used = true;
+                            Ok(docvec![
+                                "sizedInt(",
+                                value,
+                                ", ",
+                                details.size,
+                                ", ",
+                                bool(details.endianness.is_big()),
+                                ")"
+                            ])
+                        }
                     }
                 } else {
                     self.tracker.float_bit_array_segment_used = true;
@@ -319,43 +337,41 @@ impl<'module> Generator<'module> {
             .iter()
             .find(|x| matches!(x, Opt::Size { .. }));
 
-        let has_explicit_size = size.is_some();
-
-        let size = match size {
+        let (size_value, size) = match size {
             Some(Opt::Size { value: size, .. }) => {
-                let size_int = match *size.clone() {
-                    TypedExpr::Int {
-                        location: _,
-                        type_: _,
-                        value,
-                        int_value: _,
-                    } => value.parse().unwrap_or(0),
-                    _ => 0,
+                let size_value = match *size.clone() {
+                    TypedExpr::Int { int_value, .. } => Some(int_value),
+                    _ => None,
                 };
 
-                if size_int > 0 && size_int % 8 != 0 {
-                    return Err(Error::Unsupported {
-                        feature: "Non byte aligned array".into(),
-                        location: segment.location,
-                    });
+                if let Some(size_value) = size_value.as_ref() {
+                    if *size_value > BigInt::ZERO && size_value % 8 != BigInt::ZERO {
+                        return Err(Error::Unsupported {
+                            feature: "Non byte aligned array".into(),
+                            location: segment.location,
+                        });
+                    }
                 }
 
-                self.not_in_tail_position(|gen| gen.wrap_expression(size))?
+                (
+                    size_value,
+                    self.not_in_tail_position(|gen| gen.wrap_expression(size))?,
+                )
             }
             _ => {
-                let default_size = if segment.type_ == crate::type_::int() {
+                let size_value = if segment.type_ == crate::type_::int() {
                     8usize
                 } else {
                     64usize
                 };
 
-                docvec![default_size]
+                (Some(BigInt::from(size_value)), docvec![size_value])
             }
         };
 
         Ok(SizedBitArraySegmentDetails {
-            has_explicit_size,
             size,
+            size_value,
             endianness,
         })
     }
@@ -1428,19 +1444,36 @@ fn bit_array<'a>(
                 sized_bit_array_segment_details(segment, tracker, &mut constant_expr_fun)?;
 
             if segment.type_ == crate::type_::int() {
-                if details.has_explicit_size {
-                    tracker.sized_integer_segment_used = true;
-                    Ok(docvec![
-                        "sizedInt(",
-                        value,
-                        ", ",
-                        details.size,
-                        ", ",
-                        bool(details.endianness.is_big()),
-                        ")"
-                    ])
-                } else {
-                    Ok(value)
+                match (details.size_value, segment.value.as_ref()) {
+                    (Some(size_value), Constant::Int { int_value, .. })
+                        if size_value <= SAFE_INT_SEGMENT_MAX_SIZE.into() =>
+                    {
+                        let bytes = bit_array_segment_int_value_to_bytes(
+                            int_value.clone(),
+                            size_value,
+                            details.endianness,
+                            segment.location,
+                        )?;
+
+                        Ok(u8_slice(&bytes))
+                    }
+
+                    (Some(size_value), _) if size_value == 8.into() => Ok(value),
+
+                    (Some(size_value), _) if size_value <= 0.into() => Ok(docvec![]),
+
+                    _ => {
+                        tracker.sized_integer_segment_used = true;
+                        Ok(docvec![
+                            "sizedInt(",
+                            value,
+                            ", ",
+                            details.size,
+                            ", ",
+                            bool(details.endianness.is_big()),
+                            ")"
+                        ])
+                    }
                 }
             } else {
                 tracker.float_bit_array_segment_used = true;
@@ -1485,8 +1518,8 @@ fn bit_array<'a>(
 
 #[derive(Debug)]
 struct SizedBitArraySegmentDetails<'a> {
-    has_explicit_size: bool,
     size: Document<'a>,
+    size_value: Option<BigInt>, // This is set when the segment's size is known at compile time
     endianness: Endianness,
 }
 
@@ -1523,41 +1556,38 @@ fn sized_bit_array_segment_details<'a>(
         .iter()
         .find(|x| matches!(x, Opt::Size { .. }));
 
-    let has_explicit_size = size.is_some();
-
-    let size = match size {
+    let (size_value, size) = match size {
         Some(Opt::Size { value: size, .. }) => {
-            let size_int = match *size.clone() {
-                Constant::Int {
-                    location: _,
-                    value,
-                    int_value: _,
-                } => value.parse().unwrap_or(0),
-                _ => 0,
+            let size_value = match *size.clone() {
+                Constant::Int { int_value, .. } => Some(int_value),
+                _ => None,
             };
-            if size_int > 0 && size_int % 8 != 0 {
-                return Err(Error::Unsupported {
-                    feature: "Non byte aligned array".into(),
-                    location: segment.location,
-                });
+
+            if let Some(size_value) = size_value.as_ref() {
+                if *size_value > BigInt::ZERO && size_value % 8 != BigInt::ZERO {
+                    return Err(Error::Unsupported {
+                        feature: "Non byte aligned array".into(),
+                        location: segment.location,
+                    });
+                }
             }
 
-            constant_expr_fun(tracker, size)?
+            (size_value, constant_expr_fun(tracker, size)?)
         }
         _ => {
-            let default_size = if segment.type_ == crate::type_::int() {
+            let size_value = if segment.type_ == crate::type_::int() {
                 8usize
             } else {
                 64usize
             };
 
-            docvec![default_size]
+            (Some(BigInt::from(size_value)), docvec![size_value])
         }
     };
 
     Ok(SizedBitArraySegmentDetails {
-        has_explicit_size,
         size,
+        size_value,
         endianness,
     })
 }
@@ -1797,3 +1827,14 @@ fn record_constructor<'a>(
         )
     }
 }
+
+fn u8_slice<'a>(bytes: &[u8]) -> Document<'a> {
+    let s: EcoString = bytes
+        .iter()
+        .map(u8::to_string)
+        .collect::<Vec<_>>()
+        .join(", ")
+        .into();
+
+    docvec![s]
+}