Skip to main content

o1_utils/
bitwise_operations.rs

1//! This module provides a set of functions to perform bit operations on big integers.
2//! In particular, it gives XOR and NOT for `BigUint`.
3use alloc::{vec, vec::Vec};
4use core::cmp::{max, Ordering};
5use num_bigint::BigUint;
6
7use crate::BigUintHelpers;
8
9/// Bitwise operations
10pub trait BitwiseOps<Rhs = Self> {
11    /// Bitwise XOR of two `BigUint` inputs
12    fn bitwise_xor(input1: &Rhs, input: &Rhs) -> Rhs;
13
14    /// Conjunction of the bits of two `BigUint` inputs for a given number of bytes
15    fn bitwise_and(input1: &Rhs, input: &Rhs, bytes: usize) -> Rhs;
16
17    /// Negate the bits of a Self input
18    /// If it provides a larger desired `bits` than the input length then it takes the padded input of `bits` length.
19    /// Otherwise it only takes the bits of the input.
20    fn bitwise_not(input: &Rhs, bits: Option<usize>) -> Rhs;
21}
22
23impl BitwiseOps for BigUint {
24    fn bitwise_xor(input1: &Self, input2: &Self) -> Self {
25        // Pad to equal size in bytes
26        let bytes1 = input1.to_bytes_le().len();
27        let bytes2 = input2.to_bytes_le().len();
28        let in1 = to_padded_bytes(input1, bytes2);
29        let in2 = to_padded_bytes(input2, bytes1);
30        Self::from_bytes_le(
31            &in1.iter()
32                .zip(in2.iter())
33                .map(|(b1, b2)| b1 ^ b2)
34                .collect::<Vec<u8>>(),
35        )
36    }
37
38    #[allow(clippy::cast_possible_truncation)]
39    fn bitwise_not(input: &Self, bits: Option<usize>) -> Self {
40        // pad if needed / desired
41        // first get the number of bits of the input,
42        // take into account that BigUint::bits() returns 0 if the input is 0
43        let in_bits = input.bitlen();
44        let bits = max(in_bits, bits.unwrap_or(0));
45        // build vector of bits in little endian (least significant bit in position 0)
46        let mut bit_vec = vec![];
47        // negate each of the bits of the input
48        (0..bits).for_each(|i| bit_vec.push(!bit_at(input, i as u32)));
49        ToBigUint::to_biguint(&bit_vec)
50    }
51
52    fn bitwise_and(input1: &Self, input2: &Self, bytes: usize) -> Self {
53        let in1 = to_padded_bytes(input1, bytes);
54        let in2 = to_padded_bytes(input2, bytes);
55        Self::from_bytes_le(
56            &in1.iter()
57                .zip(in2.iter())
58                .map(|(b1, b2)| b1 & b2)
59                .collect::<Vec<u8>>(),
60        )
61    }
62}
63
64// Returns a BigUint as a Vec<u8> padded with zeros to a certain number of bytes
65// Panics if bytes < input.len()
66fn to_padded_bytes(input: &BigUint, bytes: usize) -> Vec<u8> {
67    let bytes_inp = input.to_bytes_le().len();
68    match bytes.cmp(&bytes_inp) {
69        Ordering::Greater => pad(input, bytes - bytes_inp),
70        Ordering::Equal | Ordering::Less => input.to_bytes_le(),
71    }
72}
73
74// Pads an input with a number of bytes
75fn pad(input: &BigUint, bytes: usize) -> Vec<u8> {
76    let mut padded = input.to_bytes_le();
77    padded.resize(bytes + padded.len(), 0u8);
78    padded
79}
80
81// Returns the bit value of a BigUint input at a certain position or zero
82fn bit_at(input: &BigUint, index: u32) -> bool {
83    if input.bit(u64::from(index)) {
84        ((input / BigUint::from(2u8).pow(index)) % BigUint::from(2u32)) == BigUint::from(1u32)
85    } else {
86        false
87    }
88}
89
90/// Converts types to a `BigUint`
91trait ToBigUint {
92    /// Converts a vector of bits in little endian to a `BigUint`
93    fn to_biguint(&self) -> BigUint;
94}
95
96impl ToBigUint for Vec<bool> {
97    fn to_biguint(&self) -> BigUint {
98        let mut bigvalue = BigUint::from(0u8);
99        let mut power = BigUint::from(1u8);
100        for bit in self {
101            bigvalue += power.clone() * BigUint::from(u8::from(*bit));
102            power *= BigUint::from(2u8);
103        }
104        bigvalue
105    }
106}