utils.zok

from "EMBED" import u32_to_bits;
from "EMBED" import u32_from_bits;
from "EMBED" import u8_to_bits;
from "EMBED" import u8_from_bits;
import "utils/casts/u8_to_field" as u8_to_field;

// Flatten a N*16 array of u32 to a 1-dimensional array of bytes
def flatten<N,F>(u32[N][16] nBy16Arr) -> u8[F] {
  // QUESTION: SHOULDN'T IT BE N*64???!!!
  u8[N*64] mut flattened = [0; N*64];
  u32 mut iTimes64 = 0;
  u32 mut iTimes64PlusJTimes4 = 0;
  for u32 i in 0..N {
    iTimes64 = 64 * i;
    for u32 j in 0..16 {
      iTimes64PlusJTimes4 = iTimes64 + 4 * j;
      // iTimes16TimesJTimes8 = iTimes16 * j * 8
      bool[32] newBits = u32_to_bits(nBy16Arr[i][j]);
      flattened[iTimes64PlusJTimes4] = u8_from_bits(newBits[0..8]);
      flattened[iTimes64PlusJTimes4+1] = u8_from_bits(newBits[8..16]);
      flattened[iTimes64PlusJTimes4+2] = u8_from_bits(newBits[16..24]);
      flattened[iTimes64PlusJTimes4+3] = u8_from_bits(newBits[24..32]);
    }
  }
  return flattened;
}

// does the opposite of flatten -- PR?
def unflatten<F,N>(u8[F] flattened) -> u32[N][16] {
  // 64 OR 128?! 
  assert(N*64 == F);
  u32[N][16] mut nBy16Arr = [[0; 16]; N];
  for u32 i in 0..N {
    for u32 j in 0..16 {
        bool[32] mut nextU32AsBits = [false; 32];
        for u32 k in 0..4 {
            bool[8] newBits = u8_to_bits(flattened[i*64+j*4+k]);
            for u32 l in 0..8 {
                nextU32AsBits[k*8+l] = newBits[l];
            }
        }
        nBy16Arr[i][j] = u32_from_bits(nextU32AsBits);
    }
  }
  return nBy16Arr;
}

def u8ArrayTou32Array<N,M>(u8[N] input) -> u32[M] {
  assert(N == 4*M);
  u32[M] mut output = [0; M];
  for u32 i in 0..M {
    output[i] = u32_from_bits([...u8_to_bits(input[i*4]), ...u8_to_bits(input[i*4+1]), ...u8_to_bits(input[i*4+2]), ...u8_to_bits(input[i*4+3])]);
  }
  return output;
}


// Returns a substring of length S, starting at fullString[idx]
// String here means a string of bytes (u8s), in any encoding
def substringAt<F,S>(u8[F] fullString, u32 idx) -> u8[S] {
	u8[S] mut r = [0; S];
  for u32 i in 0..S {
    r[i] = fullString[idx + i];
	}
  return r;
}

// Converts decimal number from a N-byte string (ASCII) to a field, e.g. string "1234" -> field 1234
// Does not allow for strings larger than 76 digits, but this can be manually changed depending on the curve
// ALTBN_128 (default curve for ZoKrates) has maximum value 21888242871839275222246405745257275088548364400416034343698204186575808495617, which has 77 digits
def decimalStringToField<N>(u8[N] string) -> field {
  assert(N < 77);
    field mut result = 0;
    for u32 i in 0..N {
        field asField = u8_to_field(string[i] - 48);
        result = result + asField * 10**(N-i-1);
    }
    return result;
}


// NOTE: needs testing?
// modified from utils/casts, TODO: PR these to utils/casts
// convert an array of bool to an array of u8
// the sizes must match (one u32 for 32 bool) otherwise an error will happen
def bool_array_to_u8_array<N, P>(bool[N] bits) -> u8[P] {
    assert(N == 8 * P);

    u8[P] mut res = [0; P];

    for u32 i in 0..P {
        res[i] = u8_from_bits(bits[8 * i..8 * (i + 1)]);
    }

    return res;
}
def u8_array_to_bool_array<N, P>(u8[N] input) -> bool[P] {
    assert(P == 8 * N);

    bool[P] mut res = [false; P];

    for u32 i in 0..N {
        bool[8] bits = u8_to_bits(input[i]);
        for u32 j in 0..8 {
            res[i * 8 + j] = bits[j];
        }
    }

    return res;
}