array.mini

/*
 * Copyright 2020, Offchain Labs, Inc. All rights reserved.
 */

type block<T> = [8]T;    // this is basically a maximum-size AVM tuple

type array<T> = struct {
    size: uint,
    topstep: uint,
    contents: block<T>,
    access: func(array<T>, uint) -> T,
};

public func array_new<T>(size: uint, base_val: T) -> array<T> {
    let chunk = 1;
    while 8*chunk < size {
	chunk = 8*chunk;
	base_val = unsafecast<T>(newfixedarray(8, base_val));
    }
    struct {
        size: size,
        topstep: chunk,
        contents: newfixedarray(8, base_val),
        access: array_get::<T>,
    }
}

public throw func array_get<T>(arr: array<T>, index: uint) -> T {
    if index >= arr.size {
	error;  // out of bounds access
    }
    let res = arr.contents;
    let chunk = arr.topstep;
    while 1 <= chunk {
	res = unsafecast<block<T> >(res[index/chunk]);
	index = index % chunk;
	chunk = chunk / 8;
    }
    unsafecast<T>(res)
}

public throw func array_get_consecutive<T>(arr: array<T>, index: uint) -> (T, T) {
    // Get two consecutive items at [index] and [index+1]
    // This will be faster than two separate arrayGets, in the common case.
    if index+1 >= arr.size {
	error;
    }
    if index % 8 == 7 {
	// leaves are not in same bottom-level block, do it the slow way
	(array_get::<T>(arr, index), array_get::<T>(arr, index+1))
    } else {
	let res = arr.contents;
	let chunk = arr.topstep;
	while 8 <= chunk {
	    res = unsafecast<block<T> >(res[index/chunk]);
	    index = index % chunk;
	    chunk = chunk / 8;
	}
	(res[index], res[index+1])
    }
}

public throw func array_set<T>(arr: array<T>, index: uint, value: T) -> array<T> {
    if index >= arr.size {
	error;
    }
    arr with { contents: arraySetInternal::<T>(arr.contents, arr.topstep, index, value) }
}

throw func arraySetInternal<T>(b: block<T>, chunkSize: uint, index: uint, value: T) -> block<T> {
    if 1 == chunkSize {
	// at a leaf of the tree
	b with {[index] = value}
    } else {
	// at an internal node of the tree
	b with {[index/chunkSize] = unsafecast<T>(arraySetInternal::<T>(
	    unsafecast<block<T> >(b[index/chunkSize]),
	    chunkSize / 8,
	    index % chunkSize,
	    value
	))}
    }
}

public throw func array_swap<T>(a: array<T>, index: uint, value: T) -> (array<T>, T) {
    // Write a new value into a slot of the array, returning the old value that was in that slot
    if index >= a.size {
	error;
    }
    let sub = arraySwapInternal::<T>(a.contents, a.topstep, index, value);
    (
	a with { contents: sub.bloc },
	sub.val,
    )
}

type arraySwapSubResult<T> = struct {
    bloc: block<T>,
    val:  T,
};

throw func arraySwapInternal<T>(b: block<T>, chunk: uint, index: uint, value: T) -> arraySwapSubResult<T> {
    if 1 == chunk {
	struct {
	    bloc: b with {[index] = value},
	    val: b[index],
	}
    } else {
        let subIndex = index/chunk;
	let sub = arraySwapInternal::<T>(
	    unsafecast<block<T> >(b[subIndex]),
	    chunk / 8,
	    index % chunk,
	    value
	);
	sub with { bloc: b with {[subIndex] = unsafecast<T>(sub.bloc) } }
    }
}

public throw func array_resize<T>(a: array<T>, newSize: uint, baseVal: T) -> array<T> {
    // Resize an array.
    // If the new size is larger, copy over all of the old contents, and fill the added elements with baseVal.
    // If the new size is smaller, copy over all of the old contents that will fit.
    
    // TODO: make this more efficient by reusing more of the tree-structure of a
    let minSize = a.size;
    if minSize > newSize {
	minSize = newSize;
    }
    let ret = array_new::<T>(newSize, baseVal);
    let i = 0;
    while i < minSize {
	ret = array_set::<T>(ret, i, array_get::<T>(a, i));
	i = i+1;
    }
    ret
}

public throw func array_get_safe<T>(arr: array<T>, index: uint) -> option<T> {
    if index >= arr.size {
	return None;
    }
    Some(array_get::<T>(arr, index))  // won't error, because we bounds-checked
}

public throw func array_set_safe<T>(arr: array<T>, index: uint, value: T) -> option<array<T> > {
    if index >= arr.size {
	return None;
    }
    Some(array_set::<T>(arr, index, value))
}

public throw func array_get_consecutive_safe<T>(arr: array<T>, index: uint) -> option<(T, T)> {
    // Return Non if access is out-of-bounds, otherwise return Some((result, result'))
    if index+1 >= arr.size {
	return None;
    }
    Some(array_get_consecutive::<T>(arr, index))
}

public throw func array_swap_safe<T>(a: array<T>, index: uint, value: T) -> option<(array<T>, T)> {
    if index >= a.size {
	return None;
    }
    let (aOut, valOut) = array_swap::<T>(a, index, value);
    Some((aOut, valOut))
}