import "@std/mem"
import "@std/arraylist"
import "@std/testing"
import "@source/ast"
import "@source/diagnostics"
import "@source/ir"
import "@source/interpreter"

@hide:file
owned_text proc(text []u8) arraylist.ArrayList(u8) ! mem.AllocError {
    bytes := arraylist.init(u8, mem.c_allocator)
    errdefer arraylist.deinit(&bytes)
    try arraylist.reserve(&bytes, text.len)
    for (text) |byte| try arraylist.append(&bytes, byte)
    return bytes
}

keeps_package_declarations_distinct_and_ir_independent test {
    diagnostic_store := diagnostics.init(mem.c_allocator)
    defer diagnostics.deinit(&diagnostic_store)
    first := try owned_text("value :: 21\nmain proc() int { return value * 2 }")
    defer arraylist.deinit(&first)
    second := try owned_text("value :: 3\nmain proc() int { return value * 2 }")
    defer arraylist.deinit(&second)
    sibling := try owned_text("helper proc() int { return 7 }")
    defer arraylist.deinit(&sibling)
    path [9]mut u8 := undefined
    memcopy!(path[..], "first.hon")
    session := init(mem.c_allocator)
    # the IR must outlive the compilation, including source buffers and interned names.
    modules [2]mut ir.Module := undefined
    {
        defer deinit(&session)
        root :: try add_package(&session)
        dependency :: try add_package(&session)
        left :: try add_source(&session, root, path[..], &first, &diagnostic_store)
        right :: try add_source(&session, dependency, "second.hon", &second, &diagnostic_store)
        third :: try add_source(&session, root, "sibling.hon", &sibling, &diagnostic_store)
        path[0] = 'X'
        try testing.expect_equal(0, first.capacity)
        try testing.expect_equal(0, first.items.len)
        try testing.expect(mem.eql(source(&session, left).path, "first.hon"))
        try testing.expect_equal(root, source(&session, left).package)
        try testing.expect_equal(dependency, source(&session, right).package)
        try testing.expect_equal(root, source(&session, third).package)
        try testing.expect_equal(left, session.packages.items[usize(root)].sources.items[0])
        try testing.expect_equal(third, session.packages.items[usize(root)].sources.items[1])
        a :: declaration(&session, left, 0)
        b :: declaration(&session, right, 0)
        try testing.expect(a.source != b.source)
        tree :: ast.view(&session.builder, source(&session, left).root)
        try testing.expect_equal(ast.name_id(&tree, a.node), ast.name_id(&tree, b.node))
        try testing.expect(mem.eql(ast.source_text(&tree, ast.statement_decl(&tree, a.node).span), "value :: 21"))
        try testing.expect(mem.eql(ast.source_text(&tree, ast.statement_decl(&tree, b.node).span), "value :: 3"))
        modules[0] = try compile(&session, root, &diagnostic_store)
        errdefer ir.deinit(&modules[0])
        modules[1] = try compile(&session, dependency, &diagnostic_store)
    }
    defer ir.deinit(&modules[0])
    defer ir.deinit(&modules[1])
    try testing.expect_equal(0, diagnostics.count(&diagnostic_store))
    left_result := try interpreter.run(mem.c_allocator, &modules[0], &diagnostic_store)
    defer interpreter.deinit(&left_result)
    right_result := try interpreter.run(mem.c_allocator, &modules[1], &diagnostic_store)
    defer interpreter.deinit(&right_result)
    try testing.expect(left_result.trap == null and right_result.trap == null)
    try testing.expect_equal(i64(42), left_result.value.data.integer)
    try testing.expect_equal(i64(6), right_result.value.data.integer)
}

recompiles_specialized_calls_without_borrowing_previous_ir test {
    diagnostic_store := diagnostics.init(mem.c_allocator)
    defer diagnostics.deinit(&diagnostic_store)
    modules [2]mut ir.Module := undefined
    {
        session := init(mem.c_allocator)
        defer deinit(&session)
        root :: try add_package(&session)
        bytes := try owned_text(
            `identity proc($T type, value T) T {
            `    if (T == i16) return value else return value.missing
            `}
            `Box proc(T type) type { return struct { value T } }
            `count proc() usize {
            `    result := 0
            `    while (result < 2) result = result + 1
            `    return result
            `}
            `Fixed proc(N usize) type {
            `    Element := i8
            `    length usize := 0
            `    while (length < N) {
            `        match length { 0: Element = i16, else: Element = i16 }
            `        length = length + 1
            `    }
            `    return match length {
            `        2: struct {
            `            values [match length { 2: length, else: usize(1) }](match true {
            `                true: Element,
            `                false: i8(128),
            `            })
            `            bias i16 = match true { true: 0, false: i8(128) }
            `        },
            `        else: i8(128),
            `    }
            `}
            `data proc() [2]i16 { return [40, 2] }
            `Tag :: enum(i16) { empty = -7, payload = 90 }
            `Choice :: union(Tag) { payload struct { values [2]i16, text string }, empty void }
            `Error :: union(enum) { failed struct { code i16, text string } }
            `Packet :: struct { choices [2]Choice, results [2](i16 ! Error) }
            `packet_data proc() Packet {
            `    return {
            `        choices = [.payload { values = [40, 2], text = "ok" }, .empty],
            `        results = [40, .failed { code = 2, text = "owned" }],
            `    }
            `}
            `Token proc($value Packet) type { return struct {} }
            `Captured proc(choice Choice) type {
            `    match choice {
            `        .payload |detail|: return struct {
            `            values [match detail.text.len { 2: usize(2), else: usize(0) }]i16
            `            bias i16 = match choice { .payload |selected|: selected.values[1], .empty: 0 }
            `        },
            `        .empty: return bool,
            `    }
            `}
            `Optional proc() type {
            `    pending ?usize := 2
            `    Result := i8
            `    while pending |length| {
            `        pending = null
            `        Result = [length]i16
            `    }
            `    return Result
            `}
            `Updated proc() type {
            `    config struct { counts [2]mut usize, bias i16 } := { counts = [0, 0], bias = 0 }
            `    original :: config
            `    for (usize(0)..2) |index| {
            `        field!(config, "counts")[index] = index + 1
            `        field!(config, "bias") = field!(config, "bias") + 21
            `    }
            `    return struct {
            `        values [config.counts[1]]i16
            `        bias i16 = config.bias + i16(original.counts[1])
            `    }
            `}
            `Collected proc() type {
            `    counts [2]mut usize := [1, 1]
            `    length usize := 0
            `    for counts |amount, index| {
            `        counts[1] = 2
            `        length = length + amount
            `        if (index == 0) continue
            `    }
            `    Element := i8
            `    for "ok" |byte| { if (byte == u8(107)) Element = i16 }
            `    return [length]Element
            `}
            `TypeArray proc() type {
            `    types [2]mut [1]mut type := [[i8], [bool]]
            `    original :: types
            `    for types |row, index| {
            `        copied := row
            `        copied[0] = u8
            `        if (row[0] == u8) return bool
            `        if (index == 0) { (types[1])[0] = i16 }
            `        if (index == 1 and row[0] != i16) return bool
            `    }
            `    if (original[1][0] != bool) return bool
            `    inferred := [[i8], [bool]]
            `    inferred = [[i16], [u8]]
            `    empty [0]type :: []
            `    for (empty) |_| return bool
            `    return struct { values [types.len]types[1][0], bias inferred[0][0] = 2 }
            `}
            `TypeRows proc() [2]mut [1]mut type {
            `    rows [2]mut [1]mut type := [[i8], [i16]]
            `    (rows[0])[0] = i16
            `    return rows
            `}
            `copy_type_rows proc(rows [2]mut [1]mut type) [2]mut [1]mut type {
            `    copied := rows
            `    (copied[0])[0] = u8
            `    if (rows[0][0] != i16) return [[bool], [bool]]
            `    return copied
            `}
            `BoundaryType proc(rows [2]mut [1]mut type) type {
            `    original :: rows
            `    changed :: copy_type_rows(rows)
            `    if (changed[0][0] != u8 or rows[0][0] != i16 or original[0][0] != i16) return bool
            `    return struct { values [rows.len]rows[0][0], bias rows[1][0] = 2 }
            `}
            `read_type_rows proc($rows [2]mut [1]mut type, value rows[0][0]) rows[1][0] {
            `    return value + rows[1][0](rows.len)
            `}
            `UninitializedRows proc() [2]mut type {
            `    types [2]mut type := undefined
            `    for (usize(0)..types.len) |index| {
            `        if (index == 0) { types[index] = struct { value i16 } }
            `        else { types[index] = types[0] }
            `    }
            `    return types
            `}
            `UninitializedType proc(types [2]mut type) type {
            `    return struct { first types[0], second types[1] }
            `}
            `UninitializedValues proc() type {
            `    rows [2]mut [2]mut usize := undefined
            `    (rows[0])[0] = 1
            `    copied := rows
            `    rows[0] = [3, 4]
            `    (copied[0])[1] = 2
            `    copied[1] = [1, 1]
            `    for (usize(0)..rows[1].len) |index| { (rows[1])[index] = index + 1 }
            `    return struct {
            `        values [copied[0][1]]i16
            `        bias i16 = i16(copied[1][0] + rows[1][0])
            `    }
            `}
            `TypeRecord proc() type {
            `    Config :: struct { child type, counts [2]mut usize, bias i16 = 2 }
            `    config Config := { child = i8, counts = [0, 0] }
            `    original :: config
            `    config.child = i16
            `    for config.counts |_, index| { config.counts[index] = index + 1 }
            `    if (original.child != i8 or original.counts[1] != 0) return bool
            `    return struct { values [config.counts[1]]config.child, bias i16 = config.bias }
            `}
            `BoundaryConfig :: struct { child type, counts [2]mut usize, bias i16 = 2 }
            `record_config proc() BoundaryConfig { return { child = i16, counts = [1, 2] } }
            `change_record_config proc(input BoundaryConfig) BoundaryConfig {
            `    copied := input
            `    copied.child = u8
            `    copied.counts[1] = 3
            `    return copied
            `}
            `RecordBoundary proc(input BoundaryConfig) type {
            `    changed :: change_record_config(input)
            `    if (input.child != i16 or input.counts[1] != 2) return bool
            `    if (changed.child != u8 or changed.counts[1] != 3) return bool
            `    return struct { values [input.counts[1]]input.child, bias i16 = input.bias }
            `}
            `RecordAlias :: alias RecordBoundary(record_config())
            `read_record_config proc($input BoundaryConfig, value input.child) input.child {
            `    return value + input.bias
            `}
            `record_pair proc() struct { type, usize } { return {i16, 2} }
            `RecordPairElement proc(input struct { type, usize }) type { return field!(input, "0") }
            `read_record_pair proc($input struct { type, usize }, value RecordPairElement(input)) RecordPairElement(input) {
            `    return value + i16(field!(input, "1"))
            `}
            `read_selected proc($choice Choice, value i16) i16 {
            `    match choice {
            `        .payload |detail|: {
            `            copied :: detail
            `            if (copied.text.len != 2) return -1
            `            return match detail.values[1] {
            `                2: match true { true: value + detail.values[1], false: value.missing },
            `                else: value.missing,
            `            }
            `        }
            `        .empty: return value.missing,
            `    }
            `}
            `main proc() i16 {
            `    if $(i16 != i16) { return $(i8(128)) }
            `    values :: $(data())
            `    box Box(i16) :: { value = identity(i16, values[0]) }
            `    fixed Fixed(count()) :: { values = [box.value, 2] }
            `    same Fixed(2) :: fixed
            `    packet :: $(packet_data())
            `    if (Token(packet_data()) != Token(packet)) return 0
            `    captured Captured(packet.choices[0]) :: { values = [19, 21] }
            `    if (captured.values[0] + captured.values[1] + captured.bias != 42) return 0
            `    optional Optional() :: captured.values
            `    if (optional[0] + optional[1] != 40) return 0
            `    updated Updated() :: { values = optional }
            `    if (updated.bias != 42 or updated.values[1] != 21) return 0
            `    collected Collected() :: [19, 21, 2]
            `    same_collected Collected() :: collected
            `    total i16 := 0
            `    for same_collected |item| { total = total + item }
            `    if (total != 42) return 0
            `    type_array TypeArray() :: { values = [19, 21] }
            `    same_type_array TypeArray() :: type_array
            `    if (same_type_array.values[0] + same_type_array.values[1] + same_type_array.bias != 42) return 0
            `    boundary BoundaryType(TypeRows()) :: { values = [19, 21] }
            `    same_boundary BoundaryType(TypeRows()) :: boundary
            `    if (same_boundary.values[0] + same_boundary.values[1] + field!(same_boundary, "bias") != 42) return 0
            `    if (copy_type_rows(TypeRows())[0][0] != u8 or copy_type_rows(TypeRows())[1][0] != i16) return 0
            `    if (TypeRows()[0][0] != i16) return 0
            `    if (read_type_rows(TypeRows(), 40) != 42 or read_type_rows(TypeRows(), 40) != 42) return 0
            `    initialized UninitializedType(UninitializedRows()) :: {
            `        first = { value = 19 }, second = { value = 21 },
            `    }
            `    same_initialized UninitializedType(UninitializedRows()) :: initialized
            `    copied_initialized := same_initialized.first
            `    initialized_values UninitializedValues() :: { values = [19, 21] }
            `    same_initialized_values UninitializedValues() :: initialized_values
            `    if (same_initialized_values.values[0] + same_initialized_values.values[1] + same_initialized_values.bias != 42) return 0
            `    type_record TypeRecord() :: { values = [19, 21] }
            `    same_type_record TypeRecord() :: type_record
            `    if (same_type_record.values[0] + same_type_record.values[1] + same_type_record.bias != 42) return 0
            `    record_boundary RecordAlias :: { values = [19, 21] }
            `    same_record_boundary RecordBoundary({ child = i16, counts = [1, 2] }) :: record_boundary
            `    if (same_record_boundary.values[0] + same_record_boundary.values[1] + same_record_boundary.bias != 42) return 0
            `    if (read_record_config(record_config(), 40) != 42) return 0
            `    if (read_record_pair(record_pair(), 40) != 42 or read_record_pair(record_pair(), 40) != 42) return 0
            `    copied_initialized = same_initialized.second
            `    payload :: packet.choices[0].payload
            `    if (payload.values[0] != 40 or payload.values[1] != 2 or payload.text.len != 2) return 0
            `    match packet.choices[1] { .empty: {}, .payload: return 0 }
            `    if ((packet.results[0] catch 0) != 40) return 0
            `    error_code :: packet.results[1] catch |err| match err {
            `        .failed |detail|: detail.code + i16(detail.text.len),
            `    }
            `    if (error_code != 7) return 0
            `    return read_selected(packet.choices[0], identity(initialized.first.value) + copied_initialized.value + same.values[$(usize(1))] + same.bias - 2)
            `}
            ,
        )
        defer arraylist.deinit(&bytes)
        _ = try add_source(&session, root, "generic.hon", &bytes, &diagnostic_store)
        modules[0] = try compile(&session, root, &diagnostic_store)
        errdefer ir.deinit(&modules[0])
        modules[1] = try compile(&session, root, &diagnostic_store)
    }
    defer ir.deinit(&modules[0])
    defer ir.deinit(&modules[1])
    try testing.expect_equal(0, diagnostics.count(&diagnostic_store))
    for (0..modules.len) |index| {
        outcome := try interpreter.run(mem.c_allocator, &modules[index], &diagnostic_store)
        defer interpreter.deinit(&outcome)
        try testing.expect(outcome.trap == null)
        try testing.expect_equal(i64(42), outcome.value.data.integer)
    }
}

dependency_traps_keep_original_spans_and_remain_demand_driven test {
    mains [2][]u8 :: [
        `dep :: import "../dependency"
        `main proc() int { return dep.fail(false) }
        `
        ,
        `dep :: import "../dependency"
        `main proc() int { return dep.fail(true) }
        `
        ,
    ]
    prefix ::
        `
        `fail proc(take bool) int {
        `    if take {
        `
    for (mains) |main, index| {
        diagnostic_store := diagnostics.init(mem.c_allocator)
        defer diagnostics.deinit(&diagnostic_store)
        module ir.Module := undefined
        dependency_source SourceId := undefined
        original diagnostics.DiagnosticId := undefined
        {
            session := init(mem.c_allocator)
            defer deinit(&session)
            root :: try add_package(&session)
            dependency :: try add_package(&session)
            root_bytes := try owned_text(main)
            defer arraylist.deinit(&root_bytes)
            dependency_bytes := try owned_text(
                `
                `fail proc(take bool) int {
                `    if take {
                `        return missing
                `    }
                `    return 42
                `}
                `
                ,
            )
            defer arraylist.deinit(&dependency_bytes)
            root_source :: try add_source(&session, root, "app/main.hon", &root_bytes, &diagnostic_store)
            dependency_source = try add_source(&session, dependency, "dependency/fail.hon", &dependency_bytes, &diagnostic_store)
            try arraylist.append(&session.imports, Import{
                node = declaration(&session, root_source, 0).node,
                target = dependency,
            })
            module = try compile(&session, root, &diagnostic_store)
        }
        defer ir.deinit(&module)
        try testing.expect_equal(1, diagnostics.count(&diagnostic_store))
        original = diagnostics.DiagnosticId(0)
        diagnostic :: diagnostics.get(&diagnostic_store, original)?
        try testing.expect_equal(diagnostics.Code.unknown_name, diagnostic.code)
        try testing.expect_equal(dependency_source, diagnostic.span.source)
        start :: u32(prefix.len + "        return ".len)
        try testing.expect_equal(start, diagnostic.span.start)
        try testing.expect_equal(start + 7, diagnostic.span.end)
        outcome := try interpreter.run(mem.c_allocator, &module, &diagnostic_store)
        defer interpreter.deinit(&outcome)
        try testing.expect_equal(1, diagnostics.count(&diagnostic_store))
        if (index == 0) {
            try testing.expect(outcome.trap == null)
            try testing.expect_equal(i64(42), outcome.value.data.integer)
        } else {
            try testing.expect_equal(original, outcome.trap?)
        }
    }
}