cache interned UInt/SInt types

Cargo.lock

@@ -319,6 +319,7 @@ dependencies = [
  "jobslot",
  "num-bigint",
  "num-traits",
+ "once_cell",
  "ordered-float",
  "petgraph",
  "serde",
@@ -521,9 +522,9 @@ dependencies = [
 
 [[package]]
 name = "once_cell"
-version = "1.19.0"
+version = "1.21.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92"
+checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
 
 [[package]]
 name = "ordered-float"

Cargo.toml

@@ -30,6 +30,7 @@ indexmap = { version = "2.5.0", features = ["serde"] }
 jobslot = "0.2.23"
 num-bigint = "0.4.6"
 num-traits = "0.2.16"
+once_cell = "1.21.3"
 ordered-float = { version = "5.1.0", features = ["serde"] }
 petgraph = "0.8.1"
 prettyplease = "0.2.20"

crates/fayalite/Cargo.toml

@@ -26,6 +26,7 @@ hashbrown.workspace = true
 jobslot.workspace = true
 num-bigint.workspace = true
 num-traits.workspace = true
+once_cell.workspace = true
 ordered-float.workspace = true
 petgraph.workspace = true
 serde_json.workspace = true

crates/fayalite/src/int.rs

@@ -10,7 +10,7 @@ use crate::{
         value_category::ValueCategoryValue,
     },
     hdl,
-    intern::{Intern, Interned, Memoize},
+    intern::{Intern, Interned, Memoize, OnceInterned},
     sim::value::{SimValue, ToSimValueWithType},
     source_location::SourceLocation,
     ty::{
@@ -65,14 +65,21 @@ pub type DynSize = ConstUsize<DYN_SIZE>;
 
 trait KnownSizeBaseSealed {}
 
-impl<const N: usize> KnownSizeBaseSealed for [(); N] {}
+impl<const N: usize> KnownSizeBaseSealed for ConstUsize<N> {}
 
 #[expect(private_bounds)]
-pub trait KnownSizeBase: KnownSizeBaseSealed {}
+pub trait KnownSizeBase: KnownSizeBaseSealed + GetInternedIntCaches {}
 
 macro_rules! impl_known_size_base {
     ($($size:literal),* $(,)?) => {
-        $(impl KnownSizeBase for [(); $size] {})*
+        $(impl KnownSizeBase for ConstUsize<$size> {})*
+        $(impl GetInternedIntCaches for ConstUsize<$size> {
+            #[inline(always)]
+            fn get_interned_int_caches() -> &'static InternedIntCaches<Self> {
+                static CACHES: InternedIntCaches<ConstUsize<$size>> = InternedIntCaches::new();
+                &CACHES
+            }
+        })*
     };
 }
 
@@ -113,12 +120,34 @@ impl_known_size_base! {
     0x200,
 }
 
+trait GetInternedIntCaches {
+    fn get_interned_int_caches() -> &'static InternedIntCaches<Self>
+    where
+        Self: KnownSize;
+}
+
+struct InternedIntCaches<Width: KnownSize> {
+    uint: OnceInterned<UIntType<Width>>,
+    sint: OnceInterned<SIntType<Width>>,
+}
+
+impl<Width: KnownSize> InternedIntCaches<Width> {
+    const fn new() -> Self {
+        Self {
+            uint: OnceInterned::new(),
+            sint: OnceInterned::new(),
+        }
+    }
+}
+
+#[expect(private_bounds)]
 pub trait KnownSize:
     GenericConstUsize
     + sealed::SizeTypeSealed
     + sealed::SizeSealed
     + Default
     + FillInDefaultedGenerics<Type = Self>
+    + GetInternedIntCaches
 {
     const SIZE: Self;
     type ArrayMatch<Element: Type>: AsRef<[Expr<Element>]>
@@ -148,7 +177,7 @@ pub trait KnownSize:
 
 impl<const N: usize> KnownSize for ConstUsize<N>
 where
-    [(); N]: KnownSizeBase,
+    ConstUsize<N>: KnownSizeBase,
 {
     const SIZE: Self = Self;
     type ArrayMatch<Element: Type> = [Expr<Element>; N];
@@ -221,6 +250,10 @@ pub trait Size:
     fn from_usize(v: usize) -> Self::SizeType {
         Self::try_from_usize(v).expect("wrong size")
     }
+    #[doc(hidden)]
+    fn interned_uint(size_type: Self::SizeType) -> Interned<UIntType<Self>>;
+    #[doc(hidden)]
+    fn interned_sint(size_type: Self::SizeType) -> Interned<SIntType<Self>>;
 }
 
 impl sealed::SizeTypeSealed for usize {}
@@ -229,6 +262,8 @@ impl SizeType for usize {
     type Size = DynSize;
 }
 
+const MAX_CACHED_INT_WIDTH: usize = 1 << 10;
+
 impl Size for DynSize {
     type ArrayMatch<Element: Type> = Box<[Expr<Element>]>;
     type ArraySimValue<Element: Type> = Box<[SimValue<Element>]>;
@@ -242,6 +277,36 @@ impl Size for DynSize {
     fn try_from_usize(v: usize) -> Option<Self::SizeType> {
         Some(v)
     }
+
+    #[doc(hidden)]
+    fn interned_uint(size_type: Self::SizeType) -> Interned<UIntType<Self>> {
+        static CACHED: [OnceInterned<UInt>; MAX_CACHED_INT_WIDTH] =
+            [const { OnceInterned::new() }; _];
+        #[cold]
+        fn intern_cold(width: usize) -> Interned<UInt> {
+            Intern::intern_sized(UInt::new(width))
+        }
+        if let Some(cached) = CACHED.get(size_type) {
+            cached.get_or_init(|| intern_cold(size_type))
+        } else {
+            intern_cold(size_type)
+        }
+    }
+
+    #[doc(hidden)]
+    fn interned_sint(size_type: Self::SizeType) -> Interned<SIntType<Self>> {
+        static CACHED: [OnceInterned<SInt>; MAX_CACHED_INT_WIDTH] =
+            [const { OnceInterned::new() }; _];
+        #[cold]
+        fn intern_cold(width: usize) -> Interned<SInt> {
+            Intern::intern_sized(SInt::new(width))
+        }
+        if let Some(cached) = CACHED.get(size_type) {
+            cached.get_or_init(|| intern_cold(size_type))
+        } else {
+            intern_cold(size_type)
+        }
+    }
 }
 
 impl<const VALUE: usize> sealed::SizeSealed for ConstUsize<VALUE> {}
@@ -267,6 +332,20 @@ impl<T: KnownSize> Size for T {
     fn try_from_usize(v: usize) -> Option<Self::SizeType> {
         if v == T::VALUE { Some(T::SIZE) } else { None }
     }
+
+    #[doc(hidden)]
+    fn interned_uint(_size_type: Self::SizeType) -> Interned<UIntType<Self>> {
+        T::get_interned_int_caches()
+            .uint
+            .get_or_init(|| UIntType::new_static().intern_sized())
+    }
+
+    #[doc(hidden)]
+    fn interned_sint(_size_type: Self::SizeType) -> Interned<SIntType<Self>> {
+        T::get_interned_int_caches()
+            .sint
+            .get_or_init(|| SIntType::new_static().intern_sized())
+    }
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]
@@ -586,7 +665,7 @@ macro_rules! impl_valueless_op_forward {
 }
 
 macro_rules! impl_int {
-    ($pretty_name:ident, $name:ident, $generic_name:ident, $value:ident, $SIGNED:literal) => {
+    ($pretty_name:ident, $name:ident, $generic_name:ident, $value:ident, $SIGNED:literal, $interned_int:ident) => {
         #[derive(Copy, Clone, PartialEq, Eq, Hash)]
         #[repr(transparent)]
         pub struct $name<Width: Size = DynSize> {
@@ -1003,7 +1082,7 @@ macro_rules! impl_int {
             type Output = $name<Width::Size>;
 
             fn index(&self, width: Width) -> &Self::Output {
-                Interned::into_inner(Intern::intern_sized($name::new(width)))
+                Interned::into_inner(Width::Size::$interned_int(width))
             }
         }
 
@@ -1184,8 +1263,22 @@ macro_rules! impl_int {
     };
 }
 
-impl_int!(UInt, UIntType, UIntWithoutGenerics, UIntValue, false);
-impl_int!(SInt, SIntType, SIntWithoutGenerics, SIntValue, true);
+impl_int!(
+    UInt,
+    UIntType,
+    UIntWithoutGenerics,
+    UIntValue,
+    false,
+    interned_uint
+);
+impl_int!(
+    SInt,
+    SIntType,
+    SIntWithoutGenerics,
+    SIntValue,
+    true,
+    interned_sint
+);
 
 impl UInt {
     /// gets the smallest `UInt` that fits `v` losslessly

crates/fayalite/src/intern.rs

@@ -4,6 +4,7 @@
 use crate::{intern::type_map::TypeIdMap, util::DefaultBuildHasher};
 use bitvec::{ptr::BitPtr, slice::BitSlice, vec::BitVec};
 use hashbrown::HashTable;
+use once_cell::race::OnceRef;
 use serde::{Deserialize, Serialize};
 use std::{
     any::{Any, TypeId},
@@ -1083,3 +1084,35 @@ pub trait Memoize: 'static + Send + Sync + Hash + Eq + Copy {
         self.get_cow(Cow::Borrowed(input))
     }
 }
+
+/// like `once_cell::race::OnceBox` but for `Interned<T>` instead of `Box<T>`
+pub struct OnceInterned<T: 'static + Send + Sync>(OnceRef<'static, T>);
+
+impl<T: 'static + Send + Sync + fmt::Debug> fmt::Debug for OnceInterned<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_tuple("OnceInterned").field(&self.get()).finish()
+    }
+}
+
+impl<T: 'static + Send + Sync> Default for OnceInterned<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<T: 'static + Send + Sync> OnceInterned<T> {
+    pub const fn new() -> Self {
+        Self(OnceRef::new())
+    }
+    pub fn set(&self, v: Interned<T>) -> Result<(), ()> {
+        self.0.set(v.inner)
+    }
+    pub fn get(&self) -> Option<Interned<T>> {
+        self.0.get().map(|inner| Interned { inner })
+    }
+    pub fn get_or_init<F: FnOnce() -> Interned<T>>(&self, f: F) -> Interned<T> {
+        Interned {
+            inner: self.0.get_or_init(|| f().inner),
+        }
+    }
+}