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