sim_lib_numbers_tensor_bit/
bit_tensor.rs

1//! Bit-packed boolean tensor storage, its `SpecTensor` backend, and the
2//! library that registers it as the `bool` element-type backend.
3
4use std::{any::Any, fmt, sync::Arc};
5
6use sim_kernel::{
7    AbiVersion, DefaultFactory, Dependency, Export, Factory, Lib, LibManifest, LibTarget, Linker,
8    Result, Symbol, Value, Version,
9};
10use sim_lib_numbers_tensor::{
11    SpecTensor, SpecTensorDescriptor, Tensor, TensorLocation, TensorStorage, checked_element_count,
12    domains, spec_tensor_descriptor_value, spec_tensor_symbol,
13};
14
15/// A boolean tensor stored as bit-packed `u64` words.
16///
17/// Each element occupies a single bit, so an `n`-element tensor is held in
18/// `ceil(n / 64)` words. The logical [`shape`](Self::shape) and element count
19/// drive layout; bitwise operations work directly on the packed words.
20#[derive(Clone)]
21pub struct BitTensor {
22    tensor: Tensor,
23}
24
25struct BitTensorStorage {
26    len: usize,
27    words: Arc<[u64]>,
28}
29
30impl BitTensor {
31    /// Packs a slice of booleans into a bit tensor of the given shape.
32    ///
33    /// Returns `None` when `bits.len()` does not match the element count
34    /// implied by `shape`.
35    pub fn from_bools(shape: Vec<usize>, bits: &[bool]) -> Option<Self> {
36        let len = checked_element_count(&shape).ok()?;
37        if len != bits.len() {
38            return None;
39        }
40        let mut words = vec![0u64; len.div_ceil(64)];
41        for (index, bit) in bits.iter().enumerate() {
42            if *bit {
43                words[index / 64] |= 1u64 << (index % 64);
44            }
45        }
46        let storage = Arc::new(BitTensorStorage {
47            len,
48            words: words.into(),
49        });
50        Tensor::from_storage(shape, domains::bool(), storage)
51            .ok()
52            .map(|tensor| Self { tensor })
53    }
54
55    /// Unpacks the tensor back into one boolean per element, in flat order.
56    pub fn to_bools(&self) -> Vec<bool> {
57        self.storage().to_bools()
58    }
59
60    /// Element-wise bitwise OR with another bit tensor of the same shape.
61    ///
62    /// Returns `None` when the shapes differ.
63    pub fn bit_or(&self, other: &Self) -> Option<Self> {
64        map_words(self, other, |left, right| left | right)
65    }
66
67    /// Element-wise bitwise XOR with another bit tensor of the same shape.
68    ///
69    /// Returns `None` when the shapes differ.
70    pub fn bit_xor(&self, other: &Self) -> Option<Self> {
71        map_words(self, other, |left, right| left ^ right)
72    }
73
74    /// Element-wise bitwise AND with another bit tensor of the same shape.
75    ///
76    /// Returns `None` when the shapes differ.
77    pub fn bit_and(&self, other: &Self) -> Option<Self> {
78        map_words(self, other, |left, right| left & right)
79    }
80
81    fn storage(&self) -> &BitTensorStorage {
82        self.tensor
83            .storage()
84            .as_any()
85            .downcast_ref::<BitTensorStorage>()
86            .expect("BitTensor must hold bit-packed storage")
87    }
88}
89
90impl BitTensorStorage {
91    fn to_bools(&self) -> Vec<bool> {
92        (0..self.len)
93            .map(|index| ((self.words[index / 64] >> (index % 64)) & 1) == 1)
94            .collect()
95    }
96}
97
98impl TensorStorage for BitTensorStorage {
99    fn dtype(&self) -> &Symbol {
100        static DTYPE: std::sync::OnceLock<Symbol> = std::sync::OnceLock::new();
101        DTYPE.get_or_init(domains::bool)
102    }
103
104    fn len(&self) -> usize {
105        self.len
106    }
107
108    fn location(&self) -> TensorLocation {
109        TensorLocation::Host
110    }
111
112    fn cell(&self, index: usize) -> Result<Value> {
113        if index >= self.len {
114            return Err(sim_kernel::Error::Eval(
115                "tensor cell index was out of bounds".to_owned(),
116            ));
117        }
118        let bit = ((self.words[index / 64] >> (index % 64)) & 1) == 1;
119        bool_value(bit)
120            .ok_or_else(|| sim_kernel::Error::Eval("bool tensor cell encode failed".to_owned()))
121    }
122
123    fn materialize(&self) -> Result<Arc<dyn TensorStorage>> {
124        Ok(Arc::new(Self {
125            len: self.len,
126            words: self.words.clone(),
127        }))
128    }
129
130    fn as_any(&self) -> &dyn Any {
131        self
132    }
133}
134
135impl fmt::Debug for BitTensor {
136    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
137        f.debug_struct("BitTensor")
138            .field("shape", &self.shape())
139            .field("bits", &self.to_bools())
140            .finish()
141    }
142}
143
144impl PartialEq for BitTensor {
145    fn eq(&self, other: &Self) -> bool {
146        self.shape() == other.shape() && self.to_bools() == other.to_bools()
147    }
148}
149
150impl Eq for BitTensor {}
151
152impl SpecTensor for BitTensor {
153    fn shape(&self) -> &[usize] {
154        self.tensor.shape()
155    }
156
157    fn dtype(&self) -> Symbol {
158        domains::bool()
159    }
160
161    fn to_uniform(&self) -> Tensor {
162        self.tensor.clone()
163    }
164
165    fn from_uniform(tensor: &Tensor) -> Option<Self> {
166        (tensor.dtype() == &domains::bool()).then_some(())?;
167        if tensor.storage().as_any().is::<BitTensorStorage>() {
168            return Some(Self {
169                tensor: tensor.clone(),
170            });
171        }
172        let bits = tensor
173            .cells()
174            .ok()?
175            .iter()
176            .map(parse_bool_cell)
177            .collect::<Option<Vec<_>>>()?;
178        Self::from_bools(tensor.shape().to_vec(), &bits)
179    }
180}
181
182/// Registered library that installs the bit-packed boolean tensor backend.
183///
184/// Loading this [`Lib`] registers a [`SpecTensor`] descriptor binding the
185/// `bool` element type to the [`BitTensor`] storage, so the base tensor domain
186/// can construct and round-trip boolean tensors through packed `u64` words.
187pub struct BitTensorLib;
188
189impl BitTensorLib {
190    /// Creates the bit-tensor library. The value is stateless; the spec-tensor
191    /// descriptor is installed when it is loaded into a
192    /// [`Cx`](sim_kernel::Cx).
193    pub fn new() -> Self {
194        Self
195    }
196}
197
198impl Default for BitTensorLib {
199    fn default() -> Self {
200        Self::new()
201    }
202}
203
204impl Lib for BitTensorLib {
205    fn manifest(&self) -> LibManifest {
206        LibManifest {
207            id: tensor_lib_symbol(),
208            version: Version(env!("CARGO_PKG_VERSION").to_owned()),
209            abi: AbiVersion { major: 0, minor: 1 },
210            target: LibTarget::HostRegistered,
211            requires: Vec::<Dependency>::new(),
212            capabilities: Vec::new(),
213            exports: vec![Export::Value {
214                symbol: tensor_spec_symbol(),
215            }],
216        }
217    }
218
219    fn load(&self, _cx: &mut sim_kernel::LoadCx, linker: &mut Linker<'_>) -> Result<()> {
220        linker.value(
221            tensor_spec_symbol(),
222            spec_tensor_descriptor_value(
223                &DefaultFactory,
224                SpecTensorDescriptor {
225                    symbol: tensor_spec_symbol(),
226                    dtype: domains::bool(),
227                    implementation: "BitTensor",
228                    storage: "canonical Tensor storage over bit-packed u64 words",
229                },
230            )?,
231        )
232    }
233}
234
235/// The manifest id symbol for this library (`numbers/tensor-bit`).
236pub fn tensor_lib_symbol() -> Symbol {
237    domains::domain("tensor-bit")
238}
239
240/// The symbol under which the bit-tensor [`SpecTensor`] descriptor is exported.
241pub fn tensor_spec_symbol() -> Symbol {
242    spec_tensor_symbol("bit")
243}
244
245fn bool_value(value: bool) -> Option<Value> {
246    DefaultFactory
247        .number_literal(domains::bool(), value.to_string())
248        .ok()
249}
250
251fn parse_bool_cell(value: &Value) -> Option<bool> {
252    let mut cx = sim_kernel::Cx::new(
253        std::sync::Arc::new(sim_kernel::NoopEvalPolicy),
254        std::sync::Arc::new(DefaultFactory),
255    );
256    let literal = value
257        .object()
258        .as_number_value()?
259        .number_literal(&mut cx)
260        .ok()??;
261    (literal.domain == domains::bool())
262        .then(|| literal.canonical.parse::<bool>().ok())
263        .flatten()
264}
265
266fn map_words(
267    left: &BitTensor,
268    right: &BitTensor,
269    f: impl Fn(u64, u64) -> u64,
270) -> Option<BitTensor> {
271    if left.shape() != right.shape() {
272        return None;
273    }
274    let storage = Arc::new(BitTensorStorage {
275        len: left.storage().len,
276        words: left
277            .storage()
278            .words
279            .iter()
280            .zip(right.storage().words.iter())
281            .map(|(left, right)| f(*left, *right))
282            .collect::<Vec<_>>()
283            .into(),
284    });
285    Tensor::from_storage(left.shape().to_vec(), domains::bool(), storage)
286        .ok()
287        .map(|tensor| BitTensor { tensor })
288}
289
290#[cfg(test)]
291mod tests {
292    use std::sync::Arc;
293
294    use sim_kernel::Lib;
295
296    use super::{BitTensor, BitTensorLib, SpecTensor, tensor_spec_symbol};
297
298    #[test]
299    fn bit_tensor_and_matches_bool_and() {
300        let left = BitTensor::from_bools(vec![4], &[true, false, true, true]).unwrap();
301        let right = BitTensor::from_bools(vec![4], &[true, true, false, true]).unwrap();
302        let out = left.bit_and(&right).unwrap();
303        assert_eq!(out.to_bools(), vec![true, false, false, true]);
304        let uniform = out.to_uniform();
305        assert_eq!(uniform.shape(), &[4]);
306    }
307
308    #[test]
309    fn uniform_roundtrip_preserves_bit_storage_identity() {
310        let tensor = BitTensor::from_bools(vec![4], &[true, false, true, true]).unwrap();
311        let uniform = tensor.to_uniform();
312        assert!(Arc::ptr_eq(tensor.tensor.storage(), uniform.storage()));
313
314        let roundtrip = BitTensor::from_uniform(&uniform).unwrap();
315        assert!(Arc::ptr_eq(roundtrip.tensor.storage(), uniform.storage()));
316        assert_eq!(roundtrip.to_bools(), vec![true, false, true, true]);
317    }
318
319    #[test]
320    fn constructor_rejects_overflowing_shape() {
321        assert!(BitTensor::from_bools(vec![usize::MAX, 2], &[]).is_none());
322    }
323
324    #[test]
325    fn lib_exports_spec_tensor_descriptor() {
326        assert_eq!(
327            BitTensorLib::new().manifest().exports[0].symbol(),
328            &tensor_spec_symbol()
329        );
330    }
331}
sim_lib_numbers_tensor_bit/bit_tensor.rs

sim_lib_numbers_tensor_bit/
bit_tensor.rs
