from "EMBED" import u32_to_bits as to_bits;
from "EMBED" import u32_from_bits as from_bits;

def right_rotate<N>(u32 e) -> u32 {
    bool[32] b = to_bits(e);
    u32 res = from_bits([...b[32-N..], ...b[..32-N]]);
    return res;
}

def main() {
	u32 e = 0x12345678;
	u32 f = 0x01234567;

	// rotate
	u32 rotated = right_rotate::<4>(e);
	assert(rotated == 0x81234567);

	// and
	assert((e & f) == 0x00204460);

	// xor
	assert((e ^ f) == 0x1317131f);

	// shift
	assert(e >> 0x0000000c == 0x00012345);
	assert(e << 0x0000000c == 0x45678000);

	// not
	assert(!e == 0xedcba987);

	// add
	assert(0xfefefefe + 0xefefefef + 0xffffffff + 0xeeeeeeee == 0xddddddda);

	// to_bits
	bool[32] bits1 = to_bits(e);
	bool[32] expected1 = [false, false, false, true, false, false, true, false, false, false, true, true, false, true,  false, false, false, true, false, true, false, true, true, false, false, true, true,  true, true, false, false, false];
	assert(bits1 == expected1);
	assert(e == from_bits(expected1));

	bool[32] bits2 = to_bits(f);
	bool[32] expected2 = [false, false, false, false, false, false, false, true, false, false, true, false, false, false, true, true, false, true, false, false, false, true,  false, true, false, true, true, false, false, true, true, true];
	assert(bits2 == expected2);
	assert(f == from_bits(expected2));

	// S0
	u32 e2 = right_rotate::<2>(e);
	u32 e13 = right_rotate::<13>(e);
	u32 e22 = right_rotate::<22>(e);
	u32 S0 = e2 ^ e13 ^ e22;
	assert(S0 == 0x66146474);

	// S1
	u32 e6 = right_rotate::<6>(e);
	u32 e11 = right_rotate::<11>(e);
	u32 e25 = right_rotate::<25>(e);
	u32 S1 = e6 ^ e11 ^ e25;
	assert(S1 == 0x3561abda);

	return;
}