midenc-hir 0.10.2

use alloc::{format, rc::Rc, string::ToString};
use core::ops::{Deref, DerefMut};

use litcheck_filecheck::{filecheck, litcheck};
use pretty_assertions::assert_eq;

use crate::{
    BuilderExt, CallConv, Context, FunctionType, Immediate, OpParser, OpRegistration, OperationRef,
    Symbol, SymbolTable, Type, UnsafeIntrusiveEntityRef, ValueRef, Visibility,
    attributes::IntegerLikeAttr,
    diagnostics::{Report, SourceSpan, Uri},
    dialects::builtin::{
        BuiltinOpBuilder, Function, Module, Ret, RetImm, UnrealizedConversionCast, World, WorldRef,
        attributes::{AbiParam, Signature},
    },
    parse::{self, ParseResult, ParserConfig},
    print::AsmPrinter,
    testing::Test,
};

type TestResult<T = ()> = Result<T, Report>;

#[test]
fn parse_simple_function() -> TestResult {
    let mut test = ParserTest::default();

    let source = "\
builtin.function public extern(\"C\") @entrypoint(%a: i32) -> i32 {
    ret %a : (i32);
};";

    let entrypoint = test.parse::<Function>("parse_simple_function.hir", source)?;
    let entrypoint = entrypoint.borrow();

    assert_eq!(entrypoint.name().as_str(), "entrypoint");
    assert_eq!(
        &*entrypoint.get_signature(),
        &Signature::new(&test.context_rc(), [Type::I32], [Type::I32])
    );
    assert_eq!(entrypoint.num_locals(), 0);
    assert_eq!(entrypoint.body().entry().body().len(), 1);

    Ok(())
}

#[test]
#[ignore]
fn parse_simple_function_generic() -> TestResult {
    let mut test = ParserTest::default();

    let source = r#""builtin.function"() <{
        name = @entrypoint,
        signature: #builtin.signature<"public extern(\"C\") (i32) -> i32">,
    }> ({
^entry(%a: i32):
    "builtin.ret" %a : (i32) -> ();
}) : () -> ();"#;

    let world = test.parse_generic("parse_simple_function_generic.hir", source)?;
    let entrypoint = world.borrow().body().entry().front().unwrap();
    let entrypoint = entrypoint.borrow();
    let entrypoint = entrypoint.downcast_ref::<Function>().expect("expected to parse a function");

    assert_eq!(entrypoint.name().as_str(), "entrypoint");
    assert_eq!(
        &*entrypoint.get_signature(),
        &Signature::new(&test.context_rc(), [Type::I32], [Type::I32])
    );
    assert_eq!(entrypoint.num_locals(), 0);
    assert_eq!(entrypoint.body().entry().body().len(), 1);

    Ok(())
}

#[test]
fn parse_module_with_intra_function_symbol_references() -> TestResult {
    let mut test = ParserTest::default();

    let source = "\
    builtin.module public @test {
        builtin.global_variable public @var : i32;

        builtin.function public extern(\"C\") @entrypoint(%a: i32) -> ptr<u8, byte> {
            %ptr = builtin.global_symbol ::@test::@var+8 : ptr<u8, byte>;
            builtin.ret %ptr : (ptr<u8, byte>);
        };
    };";

    let parsed = test.parse_any("parse_module_with_intra_function_symbol_refs.hir", source)?;
    let parsed = parsed.borrow();
    let module = parsed.downcast_ref::<Module>().unwrap();

    assert_eq!(module.get_name().as_str(), "test");
    let symbol_manager = module.symbol_manager();
    assert_eq!(symbol_manager.symbols().symbols().count(), 2);
    let var = symbol_manager
        .lookup_op("var")
        .expect("'var' was not registered in symbol table after parsing");
    let entrypoint = symbol_manager
        .lookup_op("entrypoint")
        .expect("'entrypoint' was not registered in symbol table after parsing");
    let var = var.borrow();
    let var_uses = var.as_symbol().unwrap().iter_uses().count();
    assert_eq!(var_uses, 1);

    Ok(())
}

/// The `reserved_memory` module attribute and the function-table ops are inputs to the linker's
/// memory layout, so they must survive a print/parse round-trip of the textual HIR.
#[test]
fn module_reserved_memory_and_function_table_roundtrip() -> TestResult {
    use crate::{Op, dialects::builtin::attributes::U64Attr};

    let mut test = ParserTest::default();

    let source = "\
    builtin.module public @test {
        builtin.function internal extern(\"C\") @callee(%a: i32, %b: i32) -> i32 {
            builtin.ret %a : (i32);
        };

        builtin.function_table private @tbl : 3 {
            builtin.function_table_entry 1 @callee tag 4;
        };
    };";

    let parsed = test.parse_any("function_table_roundtrip.hir", source)?;

    // Attach the linker-facing memory reservation, then round-trip through print/parse and
    // verify nothing is lost
    {
        let attr = test.context_rc().create_attribute::<U64Attr, _>(0x110000u64);
        let mut op = parsed;
        let mut op = op.borrow_mut();
        op.set_attribute(Module::RESERVED_MEMORY_ATTR, attr);
    }

    let flags = Default::default();
    let mut printer = AsmPrinter::new(test.context_rc(), &flags);
    printer.print_operation(parsed.borrow());
    let printed = printer.finish().to_string();
    std::println!("{printed}");

    let reparsed = test.parse_any("function_table_roundtrip2.hir", &printed)?;
    let reparsed = reparsed.borrow();
    let module = reparsed.downcast_ref::<Module>().unwrap();

    let reserved = module
        .as_operation()
        .get_typed_attribute::<U64Attr>(Module::RESERVED_MEMORY_ATTR)
        .expect("reserved_memory attribute lost in print/parse round-trip");
    assert_eq!(**reserved.borrow(), 0x110000);

    let table = module
        .symbol_manager()
        .lookup_op("tbl")
        .expect("'tbl' was not registered in symbol table after parsing");
    let table = table.borrow();
    let table = table
        .downcast_ref::<crate::dialects::builtin::FunctionTable>()
        .expect("expected 'tbl' to be a function table");
    assert_eq!(*table.get_num_slots(), 3);
    let entries = table.entries();
    let entries = entries.entry();
    let mut entries = entries.body().iter();
    let entry = entries.next().expect("expected one table entry");
    let entry = entry
        .downcast_ref::<crate::dialects::builtin::FunctionTableEntry>()
        .expect("expected a function table entry");
    assert_eq!(*entry.get_index(), 1);
    assert_eq!(*entry.get_type_tag(), 4, "signature tag lost in print/parse round-trip");
    assert!(entries.next().is_none());

    Ok(())
}

#[test]
fn derive_roundtrip_test() -> TestResult {
    let test = Test::new("derive_roundtrip_test", &[Type::I32], &[Type::U32]);
    let mut test = ParserTest { test };

    {
        let mut f = test.function_builder();
        let v0 = f.entry_block().borrow().arguments()[0] as ValueRef;
        let v1 = f.builder_mut().unrealized_conversion_cast(v0, Type::U32, SourceSpan::UNKNOWN)?;
        f.builder_mut().ret([v1], SourceSpan::UNKNOWN);
    }

    let flags = Default::default();
    let mut printer = AsmPrinter::new(test.context_rc(), &flags);
    printer.print_operation(test.function().borrow());
    let source = printer.render().to_string();

    let parsed = test.parse::<Function>("derive_roundtrip.hir", &source)?;
    let parsed = parsed.borrow();

    printer.print_operation(&parsed);
    let roundtripped = printer.finish().to_string();

    std::println!("{source}");
    std::println!("{roundtripped}");
    //assert_eq!(&source, &roundtripped);
    filecheck!(
        &roundtripped,
        r#"
    // CHECK: builtin.function public extern("C") @derive_roundtrip_test([[V0:%\d+]]: i32) -> u32 {
    // CHECK-NEXT: [[V1:%\d+]] = builtin.unrealized_conversion_cast [[V0]] <{ ty = #builtin.type<u32> }>;
    // CHECK-NEXT: builtin.ret [[V1]] : (u32);
    // CHECK-NEXT: };
    "#
    );

    Ok(())
}

#[test]
fn parse_ret_imm_coerces_literal_to_declared_type() -> TestResult {
    let test = ParserTest::default();

    let source = "\
builtin.function public extern(\"C\") @retconst() -> u32 {
    builtin.ret_imm 42 : u32;
};";

    let function = test.parse::<Function>("parse_ret_imm.hir", source)?;
    let printed = format!("{}", function.as_operation_ref().borrow());
    assert!(
        printed.contains("builtin.ret_imm 42 : u32"),
        "expected the declared type to survive the round trip, got:\n{printed}"
    );

    let function = function.borrow();
    let ret_imm = function
        .body()
        .entry()
        .terminator()
        .unwrap()
        .try_downcast_op::<RetImm>()
        .expect("expected the function terminator to be builtin.ret_imm");
    let imm = ret_imm.borrow().value().as_ref().as_immediate();
    assert!(
        matches!(imm, Immediate::U32(42)),
        "expected the literal to be coerced to the declared type, got {imm:?}"
    );

    // A literal that cannot be represented in the declared type must be rejected.
    let result = test.parse::<Function>(
        "parse_ret_imm_invalid.hir",
        "\
builtin.function public extern(\"C\") @retconst() -> u8 {
    builtin.ret_imm 300 : u8;
};",
    );
    assert!(result.is_err(), "expected an out-of-range immediate to be rejected");

    Ok(())
}

/// A whole `builtin.world` — the shape `midenc --emit=hir` writes, and the shape that must
/// survive being fed back in.
const WORLD_SOURCE: &str = "\
builtin.world {
    builtin.module public @lib {
        builtin.function public extern(\"C\") @main() {
            builtin.ret;
        };
    };
};";

/// The same content without the enclosing world, which the parser does have to wrap.
const MODULE_SOURCE: &str = "\
builtin.module public @lib {
    builtin.function public extern(\"C\") @main() {
        builtin.ret;
    };
};";

#[test]
fn parsing_a_world_yields_that_world_rather_than_nesting_it() -> TestResult {
    let test = ParserTest::default();

    let parsed = test.parse_any("parse_world.hir", WORLD_SOURCE)?;
    let world = parsed
        .try_downcast_op::<World>()
        .expect("expected the parsed operation to be a 'builtin.world'");

    // The world the file declares *is* the root. Nothing was created around it.
    assert!(
        parsed.parent_op().is_none(),
        "a parsed 'builtin.world' must not be nested inside another operation, but it hangs from \
         a '{}'",
        parsed.parent_op().unwrap().borrow().name()
    );

    // And because it is the root, symbols inside it have well-formed paths. A world is an
    // anonymous symbol table, and `Symbol::path` asserts that an anonymous symbol table has no
    // parent — so a world nested in a world does not merely produce a wrong path here, it
    // panics.
    let module_op = world.borrow().body().entry().front().expect("the world declares a module");
    {
        let module = module_op.borrow();
        let module = module.downcast_ref::<Module>().expect("expected 'lib' to be a module");
        let path = module.path();
        assert!(path.is_absolute(), "a symbol path is rooted at the world");
        assert_eq!(path.to_string(), "lib");
    }

    // The function's path crosses two symbol tables, so it also pins that the world contributes
    // no component of its own — it is the root, not a named ancestor.
    let function = {
        let module = module_op.borrow();
        let module = module.downcast_ref::<Module>().unwrap();
        module.body().entry().front().expect("the module defines a function")
    };
    let function = function.borrow();
    let function = function.downcast_ref::<Function>().expect("expected 'main' to be a function");
    assert_eq!(function.path().to_string(), "lib/main");

    Ok(())
}

#[test]
fn parsing_a_module_still_wraps_it_in_a_world() -> TestResult {
    let test = ParserTest::default();

    let parsed = test.parse_any("parse_module_wrapped.hir", MODULE_SOURCE)?;
    let module = parsed
        .try_downcast_op::<Module>()
        .expect("expected the parsed operation to be a 'builtin.module'");

    // The discriminating half: an operation that is not a world still gets the world it needs to
    // resolve symbols against.
    let parent = parsed
        .parent_op()
        .expect("a top-level operation that is not a world must be given a world to live in");
    assert!(
        parent.borrow().is::<World>(),
        "expected the wrapper to be a 'builtin.world', got a '{}'",
        parent.borrow().name()
    );
    assert!(parent.parent_op().is_none(), "the wrapper world must itself be the root");

    let path = module.borrow().path();
    assert!(path.is_absolute(), "a symbol path is rooted at the wrapper world");
    assert_eq!(path.to_string(), "lib");

    Ok(())
}

/// Well-formed as text, ill-formed as IR: `builtin.function` carries the `SingleRegion` trait, so
/// a function with no region at all parses (it is how a declaration is written) and then fails
/// verification with "requires exactly one region, but got 0".
const UNVERIFIABLE_SOURCE: &str = "\
builtin.module public @lib {
    builtin.function public extern(\"C\") @main();
};";

#[test]
fn parsing_verifies_what_it_parsed() {
    let test = ParserTest::default();

    // Verification used to run inside `OperationParser::finalize`, and moving it out to
    // `parse_anchored_source` left nothing pinning that it still runs at all: the tests that
    // exercise invalid IR reach the verifier by calling `recursively_verify` themselves, and the
    // one fixture that is deliberately malformed is parsed with verification turned off. Deleting
    // the call would have kept every one of them green. This is the test that would not be.
    let result = test.parse_any("parse_unverifiable.hir", UNVERIFIABLE_SOURCE);

    let err = result.err().map(|err| err.to_string()).unwrap_or_else(|| {
        panic!(
            "the default parser configuration verifies, so ill-formed IR must not parse \
             successfully"
        )
    });

    // And the failure has to be the *verifier's*: the same source parses when verification is
    // turned off, so nothing about the text itself is what rejected it.
    test.parse_any_unverified("parse_unverifiable_unverified.hir", UNVERIFIABLE_SOURCE)
        .expect("the source is well-formed as text — only verification rejects it");

    assert!(
        err.contains("invalid operation builtin.function"),
        "expected the verifier's rejection of the region-less function, got: {err}"
    );
}

#[test]
fn parse_preserves_full_width_frame_base_indices() -> TestResult {
    use crate::dialects::debuginfo::DebugInfoDialect;

    for modifier in ["local", "global"] {
        let mut test = ParserTest::default();
        test.context().get_or_register_dialect::<DebugInfoDialect>();
        let source = format!(
            r#"builtin.function public extern("C") @frame_base(%0: i32) -> i32 {{
    di.debug_declare <{{ variable = #di.variable<{{ name = "x", file = "test.rs", line = 1 }}>, expression = #di.expression<[DW_OP_fbreg({modifier}, 2147483648+0)]> }}>;
    builtin.ret %0 : (i32);
}};"#
        );

        let function = test.parse::<Function>("frame_base_full_width_index.hir", &source)?;
        let printed = format!("{}", function.as_operation_ref().borrow());
        assert!(
            printed.contains(&format!("DW_OP_fbreg({modifier}, 2147483648+0)")),
            "frame-base index was not preserved: {printed}"
        );
    }

    Ok(())
}

#[derive(Default)]
struct ParserTest {
    test: Test,
}

impl Deref for ParserTest {
    type Target = Test;

    fn deref(&self) -> &Self::Target {
        &self.test
    }
}

impl DerefMut for ParserTest {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.test
    }
}

impl ParserTest {
    #[allow(unused)]
    pub fn parse_generic(&self, name: &str, source: &str) -> TestResult<WorldRef> {
        let config = ParserConfig::new(self.test.context_rc());
        parse::parse_generic(config, Uri::new(name), source)
    }

    pub fn parse<T: OpParser + OpRegistration>(
        &self,
        name: &str,
        source: &str,
    ) -> TestResult<UnsafeIntrusiveEntityRef<T>> {
        let config = ParserConfig::new(self.test.context_rc());
        parse::parse::<T>(config, Uri::new(name), source)
    }

    pub fn parse_any(&self, name: &str, source: &str) -> TestResult<OperationRef> {
        let config = ParserConfig::new(self.test.context_rc());
        parse::parse_any(config, Uri::new(name), source)
    }

    pub fn parse_any_unverified(&self, name: &str, source: &str) -> TestResult<OperationRef> {
        let config = ParserConfig::new(self.test.context_rc()).verify_after_parse(false);
        parse::parse_any(config, Uri::new(name), source)
    }
}

/// An unregistered dialect is a different failure from an unregistered operation inside a dialect
/// that exists, and the diagnostic has to say which one it is.
#[test]
fn an_unregistered_dialect_is_named_in_the_diagnostic() {
    let test = ParserTest::default();
    let Err(err) =
        test.parse_any("unknown_dialect.hir", "builtin.module public @t { nosuchdialect.op; };")
    else {
        panic!("an unregistered dialect must not parse");
    };
    let rendered = alloc::string::ToString::to_string(&err);
    assert!(
        rendered.contains("unknown dialect") && rendered.contains("nosuchdialect"),
        "the diagnostic must name the unregistered dialect, got: {rendered}"
    );
}

/// The control: an unknown operation inside a dialect that is registered still reports as an
/// invalid operation, not as an unknown dialect.
#[test]
fn an_unknown_operation_in_a_known_dialect_still_reports_separately() {
    let test = ParserTest::default();
    let Err(err) = test.parse_any("unknown_op.hir", "builtin.module public @t { builtin.nope; };")
    else {
        panic!("an unregistered operation must not parse");
    };
    let rendered = alloc::string::ToString::to_string(&err);
    assert!(
        !rendered.contains("unknown dialect"),
        "a known dialect must not be reported as unknown, got: {rendered}"
    );
}

/// A module with an empty body defines no symbols; building its symbol table must not unwrap the
/// missing entry block.
#[test]
fn an_empty_module_body_parses() {
    let test = ParserTest::default();
    test.parse_any("empty.hir", "builtin.module public @t {};")
        .expect("an empty module body should parse");
}

/// Two symbols with one name is something the author wrote, so it must be reported rather than
/// tripping the uniqueness assertion inside the symbol table.
#[test]
fn a_duplicate_symbol_is_reported() {
    let test = ParserTest::default();
    let source = "\
builtin.module public @t {
    builtin.function internal extern(\"C\") @f(%a: i32) -> i32 {
        builtin.ret %a : (i32);
    };
    builtin.function internal extern(\"C\") @f(%b: i32) -> i32 {
        builtin.ret %b : (i32);
    };
};";
    assert!(test.parse_any("dup.hir", source).is_err(), "a duplicate symbol must not parse");
}