spl/stage2/utils/vec.spl

// Expect retcode = 0

// 类型擦除的动态数组。元素按值复制，数组只释放自身存储，不释放元素内部资源。
// Vec 拥有 raw 指向的分配；普通赋值会复制所有权，所有权转移必须使用 take。

type .{ .Error, .Kind, .Status } = @import("error.spl");
type .{ .SizeResult, .checked_add, .checked_mul, .grow_capacity, .move } = @import("memory.spl");
type .{ .alloc, .copy, .free } = @import("../sys/mem.spl");

type Vec = @self;

raw: *u8,
data: *u8,
len: usize,
cap: usize,
elem_size: usize,
elem_align: usize,

type VecResult = enum {
    Ok: Vec,
    Err: Error,
}

type PtrResult = enum {
    Ok: *u8,
    Err: Error,
}

fn blank(elem_size: usize, elem_align: usize): Vec {
    ret Vec {
        .raw = null,
        .data = null,
        .len = 0,
        .cap = 0,
        .elem_size = elem_size,
        .elem_align = elem_align,
    };
}

fn valid_alignment(alignment: usize): bool {
    ret alignment != 0 && (alignment & (alignment - 1)) == 0;
}

fn new(elem_size: usize, elem_align: usize): VecResult {
    if elem_size == 0 { ret .Err(Error.make(.InvalidArgument)); }
    if !valid_alignment(elem_align) {
        ret .Err(Error.make(.InvalidAlignment));
    }
    ret .Ok(blank(elem_size, elem_align));
}

fn valid(self: Vec): bool {
    if self.elem_size == 0 || !valid_alignment(self.elem_align) { ret false; }
    if self.len > self.cap { ret false; }
    if self.cap == 0 {
        ret self.raw == null && self.data == null && self.len == 0;
    }
    if self.raw == null || self.data == null { ret false; }
    ret (self.data as usize & (self.elem_align - 1)) == 0;
}

fn status_error(kind: Kind): Status {
    ret .Err(Error.make(kind));
}

// 计算元素区字节数。错误时不修改任何 Vec。
fn byte_count(count: usize, elem_size: usize): SizeResult {
    ret checked_mul(count, elem_size);
}

fn allocate(capacity: usize, elem_size: usize, elem_align: usize): VecResult {
    if elem_size == 0 { ret .Err(Error.make(.InvalidArgument)); }
    if !valid_alignment(elem_align) {
        ret .Err(Error.make(.InvalidAlignment));
    }
    if capacity == 0 { ret .Ok(blank(elem_size, elem_align)); }

    var bytes: usize = 0;
    match byte_count(capacity, elem_size) {
        .Ok(value) => { bytes = value; },
        .Err(error) => { ret .Err(error); },
    }

    var total: usize = 0;
    match checked_add(bytes, elem_align - 1) {
        .Ok(value) => { total = value; },
        .Err(error) => { ret .Err(error); },
    }

    var raw = alloc(total) as *u8;
    if raw == null {
        ret .Err(Error.make(.Overflow));
    }
    var address = raw as usize;
    var mask = elem_align - 1;
    if address > (~(0 as usize)) - mask {
        free(raw as *any);
        ret .Err(Error.make(.Overflow));
    }
    var aligned = (address + mask) & ~mask;
    ret .Ok(Vec {
        .raw = raw,
        .data = aligned as *u8,
        .len = 0,
        .cap = capacity,
        .elem_size = elem_size,
        .elem_align = elem_align,
    });
}

// 保证至少可容纳 need 个元素。所有检查在原对象改变前完成。
fn reserve(self: *Vec, need: usize): Status {
    if !valid(self.*) { ret status_error(.InvalidArgument); }
    if need <= self.cap { ret .Ok; }

    var new_cap = grow_capacity(self.cap, 8, need);

    var fresh: Vec = undefined;
    match allocate(new_cap, self.elem_size, self.elem_align) {
        .Ok(value) => { fresh = value; },
        .Err(error) => { ret .Err(error); },
    }

    var used_bytes: usize = 0;
    match byte_count(self.len, self.elem_size) {
        .Ok(value) => { used_bytes = value; },
        .Err(error) => {
            drop(&fresh);
            ret .Err(error);
        },
    }
    if used_bytes != 0 {
        copy(fresh.data as *any, self.data as *any, used_bytes);
    }
    fresh.len = self.len;

    var old_raw = self.raw;
    self.* = fresh;
    free(old_raw as *any);
    ret .Ok;
}

fn get(self: *Vec, index: usize): PtrResult {
    if !valid(self.*) { ret .Err(Error.make(.InvalidArgument)); }
    if index >= self.len { ret .Err(Error.make(.OutOfRange)); }
    ret .Ok(self.data + index * self.elem_size);
}

fn set(self: *Vec, index: usize, elem: *u8): Status {
    if !valid(self.*) || elem == null { ret status_error(.InvalidArgument); }
    if index >= self.len { ret status_error(.OutOfRange); }
    var dst = self.data + index * self.elem_size;
    ret move(dst[0..self.elem_size], elem[0..self.elem_size]);
}

// 追加 count 个连续元素。elems 可以指向 self 当前已初始化的元素区。
fn append(self: *Vec, elems: *u8, count: usize): Status {
    if !valid(self.*) { ret status_error(.InvalidArgument); }
    if count == 0 { ret .Ok; }
    if elems == null { ret status_error(.InvalidArgument); }

    var copy_bytes: usize = 0;
    match byte_count(count, self.elem_size) {
        .Ok(value) => { copy_bytes = value; },
        .Err(error) => { ret .Err(error); },
    }
    var new_len: usize = 0;
    match checked_add(self.len, count) {
        .Ok(value) => { new_len = value; },
        .Err(error) => { ret .Err(error); },
    }

    var used_bytes: usize = 0;
    match byte_count(self.len, self.elem_size) {
        .Ok(value) => { used_bytes = value; },
        .Err(error) => { ret .Err(error); },
    }

    var source_inside = false;
    var source_offset: usize = 0;
    if used_bytes != 0 {
        var base = self.data as usize;
        var source = elems as usize;
        if source >= base && source - base < used_bytes {
            source_offset = source - base;
            if copy_bytes > used_bytes - source_offset {
                ret status_error(.OutOfRange);
            }
            source_inside = true;
        }
    }

    match reserve(self, new_len) {
        .Ok => {},
        .Err(error) => { ret .Err(error); },
    }
    if source_inside { elems = self.data + source_offset; }

    var dst = self.data + used_bytes;
    match move(dst[0..copy_bytes], elems[0..copy_bytes]) {
        .Ok => {},
        .Err(error) => { ret .Err(error); },
    }
    self.len = new_len;
    ret .Ok;
}

fn push(self: *Vec, elem: *u8): Status {
    ret append(self, elem, 1);
}

// 将末元素复制到 out，然后缩短数组。out 可以与元素区重叠。
fn pop(self: *Vec, out: *u8): Status {
    if !valid(self.*) || out == null { ret status_error(.InvalidArgument); }
    if self.len == 0 { ret status_error(.OutOfRange); }
    var source = self.data + (self.len - 1) * self.elem_size;
    match move(out[0..self.elem_size], source[0..self.elem_size]) {
        .Ok => {},
        .Err(error) => { ret .Err(error); },
    }
    self.len = self.len - 1;
    ret .Ok;
}

fn clear(self: *Vec) {
    self.len = 0;
}

// 成功时先释放 out 原有存储再提交副本；失败保持 out 不变。
// out 与 src 指向同一个 Vec 也允许。
fn clone(out: *Vec, src: Vec): Status {
    if !valid(out.*) || !valid(src) { ret status_error(.InvalidArgument); }
    if out.elem_size != src.elem_size || out.elem_align != src.elem_align {
        ret status_error(.InvalidArgument);
    }

    var fresh = blank(src.elem_size, src.elem_align);
    match reserve(&fresh, src.len) {
        .Ok => {},
        .Err(error) => { ret .Err(error); },
    }
    var bytes: usize = 0;
    match byte_count(src.len, src.elem_size) {
        .Ok(value) => { bytes = value; },
        .Err(error) => {
            drop(&fresh);
            ret .Err(error);
        },
    }
    if bytes != 0 {
        copy(fresh.data as *any, src.data as *any, bytes);
    }
    fresh.len = src.len;
    drop(out);
    out.* = fresh;
    ret .Ok;
}

// 显式转移所有权，返回值必须由接收方最终 drop。
fn take(source: *Vec): Vec {
    var result = source.*;
    source.* = blank(result.elem_size, result.elem_align);
    ret result;
}

fn drop(self: *Vec) {
    var elem_size = self.elem_size;
    var elem_align = self.elem_align;
    free(self.raw as *any);
    self.* = blank(elem_size, elem_align);
}

fn expect_ok(result: Status) {
    match result {
        .Ok => {},
        .Err(error) => { @assert(false, "unexpected Vec error"); },
    }
}

fn expect_error(result: Status, overflow: bool) {
    match result {
        .Ok => { @assert(false, "expected Vec error"); },
        .Err(error) => {
            match error.kind {
                .Overflow => { @assert(overflow, "unexpected overflow"); },
                .OutOfRange => { @assert(!overflow, "unexpected range error"); },
                _ => { @assert(false, "unexpected Vec error kind"); },
            }
        },
    }
}

fn make_vec(elem_size: usize, elem_align: usize): Vec {
    match new(elem_size, elem_align) {
        .Ok(value) => { ret value; },
        .Err(error) => { @assert(false, "Vec construction"); },
    }
    ret blank(1, 1);
}

fn expect_i32(self: *Vec, index: usize, expected: i32) {
    match get(self, index) {
        .Ok(ptr) => { @assert((ptr as *i32).* == expected, "Vec element value"); },
        .Err(error) => { @assert(false, "Vec element access"); },
    }
}

fn test_construction_and_alignment() {
    match new(0, 1) {
        .Ok(value) => { @assert(false, "zero element size"); },
        .Err(error) => {
            match error.kind { .InvalidArgument => {}, _ => { @assert(false); }, }
        },
    }
    match new(1, 3) {
        .Ok(value) => { @assert(false, "invalid element alignment"); },
        .Err(error) => {
            match error.kind { .InvalidAlignment => {}, _ => { @assert(false); }, }
        },
    }

    var alignment: usize = 1;
    while alignment <= 64 {
        var v = make_vec(3, alignment);
        expect_ok(reserve(&v, 9));
        @assert((v.data as usize & (alignment - 1)) == 0, "requested alignment");
        @assert(v.cap >= 9 && v.len == 0);
        drop(&v);
        @assert(v.raw == null && v.data == null && v.cap == 0);
        alignment = alignment * 2;
    }
}

fn test_growth_and_access() {
    var v = make_vec(@sizeof(i32), @alignof(i32));
    defer { drop(&v); }
    var i: i32 = 0;
    while i < 1000 {
        var value = i * 3 - 7;
        expect_ok(push(&v, &value as *u8));
        i = i + 1;
    }
    @assert(v.len == 1000 && v.cap >= v.len);
    i = 0;
    while i < 1000 {
        expect_i32(&v, i as usize, i * 3 - 7);
        i = i + 1;
    }

    var replacement: i32 = 123456;
    expect_ok(set(&v, 500, &replacement as *u8));
    expect_i32(&v, 500, replacement);
    expect_error(set(&v, 1000, &replacement as *u8), false);

    var popped: i32 = 0;
    expect_ok(pop(&v, &popped as *u8));
    @assert(popped == 2990 && v.len == 999, "pop copies then removes");
}

fn test_self_reference() {
    var v = make_vec(@sizeof(i32), @alignof(i32));
    defer { drop(&v); }
    var i: i32 = 0;
    while i < 8 {
        expect_ok(push(&v, &i as *u8));
        i = i + 1;
    }
    @assert(v.len == 8 && v.cap == 8, "prepare growth boundary");

    // 源指向旧分配，append 必须在扩容后重新定位它。
    expect_ok(append(&v, v.data, v.len));
    @assert(v.len == 16);
    i = 0;
    while i < 16 {
        expect_i32(&v, i as usize, i % 8);
        i = i + 1;
    }

    var preserved_len = v.len;
    expect_error(append(&v, v.data + (v.len - 1) * v.elem_size, 2), false);
    @assert(v.len == preserved_len, "invalid self range preserves Vec");
    expect_i32(&v, 15, 7);

    var last_ptr: *u8 = null;
    match get(&v, 15) {
        .Ok(ptr) => { last_ptr = ptr; },
        .Err(error) => { @assert(false); },
    }
    expect_ok(set(&v, 0, last_ptr));
    expect_i32(&v, 0, 7);
}

fn test_clone_take_and_failures() {
    var source = make_vec(@sizeof(i32), @alignof(i32));
    defer { drop(&source); }
    var one: i32 = 1;
    var two: i32 = 2;
    expect_ok(push(&source, &one as *u8));
    expect_ok(push(&source, &two as *u8));

    var copy = make_vec(@sizeof(i32), @alignof(i32));
    defer { drop(&copy); }
    expect_ok(clone(&copy, source));
    expect_i32(&copy, 0, 1);
    one = 99;
    expect_ok(set(&source, 0, &one as *u8));
    expect_i32(&copy, 0, 1);

    // out 与 src 为同一对象仍产生内容相同的新分配。
    var old_data = copy.data;
    expect_ok(clone(&copy, copy));
    @assert(copy.data != old_data, "self clone replaces allocation");
    expect_i32(&copy, 1, 2);

    var moved = take(&copy);
    defer { drop(&moved); }
    @assert(copy.raw == null && copy.data == null && copy.len == 0 && copy.cap == 0);
    expect_i32(&moved, 0, 1);

    var old_raw = moved.raw;
    var old_len = moved.len;
    var old_cap = moved.cap;
    expect_error(reserve(&moved, (~(0 as usize)) / moved.elem_size + 1), true);
    @assert(moved.raw == old_raw && moved.len == old_len && moved.cap == old_cap,
            "reserve failure preserves Vec");

    clear(&moved);
    @assert(moved.len == 0 && moved.cap == old_cap, "clear retains capacity");
    var output: i32 = 55;
    expect_error(pop(&moved, &output as *u8), false);
    @assert(output == 55, "failed pop preserves output");
}

fn main() {
    test_construction_and_alignment();
    test_growth_and_access();
    test_self_reference();
    test_clone_take_and_failures();
}