jitir.py

# Copyright 2025 Can Joshua Lehmann
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import code

from lwir import *

class CountVarargsValueType(BaseType):
    def __init__(self):
        super().__init__()

    def is_value(self):
        return True

    def format(self, ir):
        return "size_t"

    def __hash__(self):
        return hash("CountVarargsValueType")
    
    def __eq__(self, other):
        return isinstance(other, CountVarargsValueType)

    def __repr__(self):
        return "CountVarargsValueType()"

class LLVMAPIPlugin:
    def __init__(self, prefix, type_substitutions):
        self.prefix = prefix
        self.type_substitutions = type_substitutions

    def run(self, ir):
        defs = ""
        for inst in ir.insts:
            defs += f"  llvm::FunctionCallee {inst.format_builder_name(ir)};\n"
        
        inits = ""
        for inst in ir.insts:
            arg_types = ["llvm::PointerType::get(context, 0)"] # Builder
            for arg in inst.args:
                arg_types.append(self.type_substitutions[arg.type])
            arg_types = ", ".join(arg_types)
            inits += f"    {inst.format_builder_name(ir)} = module->getOrInsertFunction(\n"
            inits += f"      \"{self.prefix}_{inst.format_builder_name(ir)}\",\n"
            inits += f"      llvm::FunctionType::get(\n"
            inits += f"        llvm::PointerType::get(context, 0),\n"
            inits += f"        std::vector<llvm::Type*>({{ {arg_types} }}),\n"
            inits += f"        false\n"
            inits += f"      )\n"
            inits += f"    );\n"
        
        return {
            "llvmapi_defs": defs,
            "llvmapi_inits": inits
        }

class BuildBuildInstPlugin:
    def __init__(self, type_substitutions):
        self.type_substitutions = type_substitutions

    def run(self, ir):
        code = f"inline llvm::Value* build_build_inst(llvm::IRBuilder<>& builder,\n"
        code += f"                                    LLVM_API& api,\n"
        code += f"                                    Inst* inst,\n"
        code += f"                                    llvm::Value* jitir_builder,\n"
        code += f"                                    const std::vector<llvm::Value*> args) {{\n"
        code += f"  llvm::LLVMContext& context = builder.GetInsertBlock()->getModule()->getContext();\n"
        for inst in ir.insts:
            name = inst.format_name(ir)
            code += f"  if ({name}* i = dynamic_cast<{name}*>(inst)) {{\n"
            code += f"    std::vector<llvm::Value*> build_args;\n"
            code += f"    build_args.push_back(jitir_builder);\n"

            value_index = 0
            varargs = None
            for arg in inst.args:
                match arg.type:
                    case CountVarargsValueType():
                        code += f"    build_args.push_back(llvm::ConstantInt::get({self.type_substitutions[arg.type]}, {arg.name}.size() - {value_index}, false));\n"
                        varargs = arg
                    case ValueType():
                        code += f"    build_args.push_back(args[{value_index}]);\n"
                        value_index += 1
                    case Type():
                        llvm_type = self.type_substitutions[arg.type]
                        if "PointerType" in llvm_type:
                            code += f"    build_args.push_back(builder.CreateIntToPtr(llvm::ConstantInt::get(llvm::Type::getInt64Ty(context), (uint64_t)i->{arg.name}(), false), {llvm_type}));\n"
                        else:
                            code += f"    build_args.push_back(llvm::ConstantInt::get({llvm_type}, (uint64_t)i->{arg.name}(), false));\n"
            code += f"    assert(build_args.size() == {len(inst.args) + 1});\n"
            code += f"    llvm::Value* built = builder.CreateCall(api.{inst.format_builder_name(ir)}, build_args);\n"

            if varargs is not None:
                code += f"    for (size_t it = {value_index}; it < args.size(); it++) {{\n"
                code += f"      builder.CreateCall(api.set_arg, {{\n"
                code += f"        built,\n"
                code += f"        llvm::ConstantInt::get(llvm::Type::getInt64Ty(context), it, false),\n"
                code += f"        args[it]\n"
                code += f"      }});\n"
                code += f"    }}\n"

            code += f"    return built;\n"
            code += f"  }}\n"
        code += f"  assert(false && \"Unknown instruction\");\n"
        code += f"  return nullptr;\n"
        code += f"}}\n"
        return {"build_build_inst": code}

class MapSymbolsInstPlugin:
    def __init__(self, prefix):
        self.prefix = prefix

    def run(self, ir):
        code = ""
        for inst in ir.insts:
            name = f"{self.prefix}_{inst.format_builder_name(ir)}"
            code += f"map_symbol({name});\n"
        return {"map_symbols": code}

class AllocatorBuilderPlugin(BuilderPlugin):
    def allocator(self, inst, ir):
        name = inst.format_name(ir)
        arg_count = 0
        varargs = None
        for arg in inst.args:
            if isinstance(arg.type, CountVarargsValueType):
                varargs = arg
            elif arg.type.is_value():
                arg_count += 1
        
        if varargs is None:
            arg_count = str(arg_count)
        else:
            arg_count = f"({arg_count} + {varargs.name})"
        
        size = f"sizeof({name}) + sizeof(Value*) * {arg_count}"
        return f"({name}*) _section->allocator().alloc({size}, alignof({name}))"

class InstTrailingConstructorPlugin:
    def run(self, inst, ir):
        name = inst.format_name(ir)
        base = inst.base or ir.inst_base

        init_list = []
        arg_init = ""
        arg_count = 0
        varargs = None
        for arg in inst.args:
            match arg.type:
                case CountVarargsValueType():
                    varargs = arg
                case ValueType():
                    arg_init += f".with({arg_count}, {arg.name})"
                    arg_count += 1
                case Type():
                    init_list.append(f"_{arg.name}({arg.name})")
                case _:
                    assert False, f"Unknown type: {arg.type}"
        
        if varargs is None:
            arg_count = str(arg_count)
        else:
            arg_count = f"({arg_count} + {varargs.name})"
            arg_init = f".zeroed()" + arg_init

        arg_init = f"lwir::Span<Value*>::trailing(this, {arg_count})" + arg_init

        init_list = [f"{base}({inst.type}, {arg_init})"] + init_list
        init_list = ", ".join(init_list)
        
        ctor_args = ", ".join(inst.format_formal_args(ir))
        code = f"  {name}({ctor_args}): {init_list} {{\n"
        for check in inst.type_checks:
            code += f"    assert({check});\n"
        code += f"  }}\n"
        code += "\n"

        return code


class InstReadPlugin:
    def run(self, ir):
        code = "Value* read_opcode(std::string opcode) {\n"
        else_prefix = ""
        for inst in ir.insts:
            code += f"  {else_prefix}if (opcode == \"{inst.name}\") {{\n"
            # read arguments
            need_comma = False
            varargs = None
            build_args = []
            for arg in inst.args:
                if need_comma:
                    code += "    expect_char(','); skip_whitespace();\n"
                need_comma = True
                if arg.type == ValueType():
                    code += f"    Value* {arg.name} = read_value_arg();\n"
                    build_args.append(arg.name)
                    continue
                elif arg.type == CountVarargsValueType():
                    code += f"    std::vector<Value*> {arg.name} = read_value_arg_list();\n"
                    need_comma = False # read_value_arg_list needs to consume the comma
                    varargs = arg.name
                    build_args.append(f"{arg.name}.size()")
                    continue
                build_args.append(arg.name)
                code += f'    expect_word("{arg.name}");\n'
                code += f"    expect_char('=');\n"
                if arg.type == Type("Block*"):
                    code += f"    Block* {arg.name} = read_block_argument();\n"
                elif arg.type == Type("Type"):
                    code += f"    Type {arg.name} = read_type();\n"
                elif arg.type == Type("CallConv"):
                    code += f"    CallConv {arg.name} = read_call_conv();\n"
                elif arg.type == Type("LoadFlags"):
                    code += f"    LoadFlags {arg.name} = read_load_flags();\n"
                elif arg.type == Type("AliasingGroup") or arg.type == Type("uint64_t"):
                    code += f"    uint64_t {arg.name} = read_uint64();\n"
                elif arg.type == Type("const char*"):
                    code += f'    error("const char* argument not supported yet"); char * {arg.name};\n'
                else:
                    import pdb;pdb.set_trace()
                    assert False, f"Unknown argument type: {arg.type}"
            build_args = ", ".join(build_args)
            builder_call = f"_builder.{inst.format_builder_name(ir)}({build_args})\n"
            if varargs is not None:
                code += f"    Inst* result = {builder_call};\n"
                code += f"    for (size_t it = {0}; it < {varargs}.size(); it++) {{\n"
                code += f"      result->set_arg(it, {varargs}[it]);\n"
                code += f"    }}\n"
                code += f"    return result;"
            else:
                code += f"    return {builder_call};\n"
            else_prefix = "} else "
        code += '  } else { error("unknown operation " + opcode); return nullptr; }\n'
        code += "}"
        return dict(read_opcode=code)


class PrettyInstWritePlugin(InstWritePlugin):
    def __init__(self):
        super().__init__(
            overrides={
                "Comment": lambda stream: "stream << Highlight::Comment << \"; \" << _text << Highlight::None;"
            },
            stream_type="metajit::PrettyStream&"
        )
    
    def write_name(self, inst, stream):
        return f"{stream} << Highlight::Keyword << \"{inst.name}\" << Highlight::None;"

    def write_arg(self, arg, value, stream):
        code = f"{stream} << Highlight::ArgName << \"{arg.name}=\" << Highlight::None; "
        if arg.type == Type("Block*"):
            code += f"{value}->write_arg({stream});"
        elif arg.type == Type("Type") or arg.type == Type("CallConv"):
            code += f"{stream} << Highlight::Type << {value} << Highlight::None;"
        elif arg.type == Type("LoadFlags"):
            code += f"{stream} << {value};"
        elif arg.type == Type("AliasingGroup") or arg.type == Type("uint64_t"):
            code += f"{stream} << Highlight::Constant << {value} << Highlight::None;"
        else:
            assert False
        return code

class JitirInstWriteJsonPlugin(InstWriteJsonPlugin):
    def __init__(self):
        super().__init__()
    
    def write_misc(self, stream):
        return {
            "name": f"{stream} << name()",
            "type": f"{stream} << \"\\\"\" << type() << \"\\\"\""
        }

    def write_arg(self, arg, value, stream):
        if arg.type == Type("Block*"):
            return f"{stream} << {value}->name()"
        elif arg.type == Type("Type") or arg.type == Type("CallConv"):
            return f"{stream} << \"\\\"\" << {value} << \"\\\"\""
        elif arg.type == Type("LoadFlags"):
            return f"{value}.write_json({stream})"
        elif arg.type == Type("AliasingGroup") or arg.type == Type("uint64_t"):
            return f"{stream} << {value}"
        elif arg.type == Type("const char*"):
            return f"{stream} << \"\\\"\" << escape_json({value}) << \"\\\"\""
        else:
            assert False

class ClonePlugin:
    def run(self, ir):
        code = ""
        for inst in ir.insts:
            name = inst.format_name(ir)
            code += f"if (dynamic_cast<{name}*>(inst)) {{\n"
            code += f"  {name}* clone = _builder.{inst.format_builder_name(ir)}("
            args = []
            value_index = 0
            varargs = None
            for arg in inst.args:
                match arg.type:
                    case CountVarargsValueType():
                        args.append(f"inst->arg_count() - {value_index}")
                        varargs = arg
                    case ValueType():
                        args.append(f"clone_arg(inst->arg({value_index}))")
                        value_index += 1
                    case Type(name = "Block*"):
                        args.append(f"_blocks[(({name}*) inst)->{arg.name}()->name()]")
                    case Type():
                        args.append(f"(({name}*) inst)->{arg.name}()")
            code += ", ".join(args)
            code += f");\n"

            if varargs is not None:
                code += f"  for (size_t it = {value_index}; it < inst->arg_count(); it++) {{\n"
                code += f"    clone->set_arg(it, clone_arg(inst->arg(it)));\n"
                code += f"  }}\n"

            code += f"  return clone;\n"
            code += f"}}\n"
        code += f"assert(false && \"Unknown instruction\");\n"
        code += f"return nullptr;\n"
        return {"clone": code}

def binop(name, type_checks = None):
    if type_checks is None:
        type_checks = ["is_int(a->type())"]
    return Inst(name,
        args = [Arg("a"), Arg("b")],
        type = "a->type()",
        type_checks = [
            "a->type() == b->type()",
            *type_checks
        ]
    )

def binop_f(name):
    return binop(name, type_checks = ["is_float(a->type())"])

def cmp(name, type_checks):
    return Inst(name,
        args = [Arg("a"), Arg("b")],
        type = "Type::Bool",
        type_checks = [
            "a->type() == b->type()",
            *type_checks
        ]
    )

jitir = IR(
    insts = [
        Inst("Freeze",
            args = [Arg("a")],
            type = "a->type()",
            type_checks = [],
            doc = "Bind a value at tracing time. Insert a guard to check that the value is the same at runtime."
        ),
        Inst("AssumeConst",
            args = [Arg("a")],
            type = "a->type()",
            type_checks = [],
            doc = "Bind a value at tracing time without inserting a guard."
        ),
        Inst("Select",
            args = [Arg("cond", getter=Getter.Always), Arg("a"), Arg("b")],
            type = "a->type()",
            type_checks = [
                "cond->type() == Type::Bool",
                "a->type() == b->type()"
            ]
        ),
        Inst("ResizeU",
            args = [Arg("a"), Arg("type", Type("Type"))],
            type = "type",
            type_checks = [
                "is_int_or_bool(a->type())",
                "is_int_or_bool(type)"
            ],
            doc = "Zero-extend or truncate value."
        ),
        Inst("ResizeS",
            args = [Arg("a"), Arg("type", Type("Type"))],
            type = "type",
            type_checks = [
                "is_int_or_bool(a->type())",
                "is_int_or_bool(type)"
            ],
            doc = "Sign-extend or truncate value."
        ),
        Inst("ResizeX",
            args = [Arg("a"), Arg("type", Type("Type"))],
            type = "type",
            type_checks = [
                "is_int_or_bool(a->type())",
                "is_int_or_bool(type)"
            ],
            doc = "Extend or truncate value. If extending, the new bits are undefined."
        ),
        Inst("IntToFloatS",
            args = [Arg("a"), Arg("type", Type("Type"))],
            type = "type",
            type_checks = [
                "is_int(a->type())",
                "is_float(type)"
            ]
        ),
        Inst("FloatToIntS",
            args = [Arg("a"), Arg("type", Type("Type"))],
            type = "type",
            type_checks = [
                "is_float(a->type())",
                "is_int(type)"
            ]
        ),
        Inst("Load",
            args = [
                Arg("ptr", getter=Getter.Always),
                Arg("type", Type("Type")),
                Arg("flags", Type("LoadFlags"), setter=True),
                Arg("aliasing", Type("AliasingGroup"), setter=True),
                Arg("offset", Type("uint64_t"), setter=True)
            ],
            type = "type",
            type_checks = ["ptr->type() == Type::Ptr"]
        ),
        Inst("Store",
            args = [
                Arg("ptr", getter=Getter.Always),
                Arg("value", getter=Getter.Always),
                Arg("aliasing", Type("AliasingGroup"), setter=True),
                Arg("offset", Type("uint64_t"), setter=True)
            ],
            type = "Type::Void",
            type_checks = ["ptr->type() == Type::Ptr"]
        ),
        Inst("Alloca",
            args = [
                Arg("size", getter=Getter.Always)
            ],
            type = "Type::Ptr",
            type_checks = [
                "size->type() == Type::Int64"
            ],
            doc = "Allocate untyped memory on the stack."
        ),
        Inst("AddPtr",
            args = [
                Arg("ptr", getter=Getter.Always),
                Arg("offset", getter=Getter.Always)
            ],
            type = "Type::Ptr",
            type_checks = [
                "ptr->type() == Type::Ptr",
                "offset->type() == Type::Int64"
            ]
        ),
        binop("Add"),
        binop("Sub"),
        binop("Mul"),
        binop("DivS"),
        binop("DivU"),
        binop("ModS"),
        binop("ModU"),
        binop("And", type_checks = ["is_int_or_bool(a->type())"]),
        binop("Or", type_checks = ["is_int_or_bool(a->type())"]),
        binop("Xor", type_checks = ["is_int_or_bool(a->type())"]),
        binop("ShrU"),
        binop("ShrS"),
        binop("Shl"),
        binop_f("AddF"),
        binop_f("SubF"),
        binop_f("MulF"),
        binop_f("DivF"),
        cmp("Eq", type_checks = []),
        cmp("LtU", type_checks = ["is_int(a->type())"]),
        cmp("LtS", type_checks = ["is_int(a->type())"]),
        cmp("LtFO", type_checks = ["is_float(a->type())"]),
        cmp("LtFU", type_checks = ["is_float(a->type())"]),
        Inst("Call",
            args = [
                Arg("callee", getter=Getter.Always),
                Arg("args", type=CountVarargsValueType()),
                Arg("type", Type("Type")),
                Arg("call_conv", Type("CallConv"), setter=True),
            ],
            type = "type",
            type_checks = [
                "callee->type() == Type::Ptr"
            ],
            doc = "Call a function pointer."
        ),
        Inst("Branch",
            args = [
                Arg("cond", getter=Getter.Always),
                Arg("true_block", type=Type("Block*"), setter=True),
                Arg("false_block", type=Type("Block*"), setter=True)
            ],
            type = "Type::Void",
            type_checks = ["cond->type() == Type::Bool"],
            doc = "Conditional jump."
        ),
        Inst("Jump",
            args = [
                Arg("args", type=CountVarargsValueType()),
                Arg("block", type=Type("Block*"), setter=True)
            ],
            type = "Type::Void",
            type_checks = [],
            doc = "Unconditional jump."
        ),
        Inst("Exit",
            args = [],
            type = "Type::Void",
            type_checks = [],
            doc = "Return from section."
        ),
        Inst("Comment",
            args = [Arg("text", Type("const char*"), setter=True)],
            type = "Type::Void",
            type_checks = [],
            doc = "No-op. Used for adding comments to the IR."
        ),
    ]
)

lwir(
    template_path = "jitir.tmpl.hpp",
    output_path = "jitir.hpp",
    ir = jitir,
    plugins = [
        InstPlugin([
            InstTrailingConstructorPlugin(),
            InstGetterPlugin(),
            InstSetterPlugin(),
            PrettyInstWritePlugin(),
            JitirInstWriteJsonPlugin(),
            InstEqualsPlugin(),
            InstHashPlugin()
        ]),
        AllocatorBuilderPlugin(),
        ClonePlugin(),
        CAPIPlugin(
            prefix = "jitir",
            builder_name = "::metajit::TraceBuilder",
            type_substitutions = {
                Type("size_t"): "uint64_t",
                Type("uint64_t"): "uint64_t",
                Type("Type"): "uint32_t",
                Type("CallConv"): "uint32_t",
                Type("LoadFlags"): "uint32_t",
                Type("InputFlags"): "uint32_t",
                Type("Block*"): "void*",
                Type("AliasingGroup"): "uint32_t", # Needs to be passed by LLVM IR
                Type("const char*"): "const char*",
                ValueType(): "void*",
                CountVarargsValueType(): "uint64_t"
            }
        ),
        InstReadPlugin(),
    ],
    placeholder = lambda name: "/* ${" + name + "} */"
)

llvm_type_substitutions = {
    Type("size_t"): "llvm::Type::getInt64Ty(context)",
    Type("uint64_t"): "llvm::Type::getInt64Ty(context)",
    Type("Type"): "llvm::Type::getInt32Ty(context)",
    Type("CallConv"): "llvm::Type::getInt32Ty(context)",
    Type("LoadFlags"): "llvm::Type::getInt32Ty(context)",
    Type("InputFlags"): "llvm::Type::getInt32Ty(context)",
    Type("Block*"): "llvm::PointerType::get(context, 0)",
    Type("AliasingGroup"): "llvm::Type::getInt32Ty(context)",
    Type("const char*"): "llvm::PointerType::get(context, 0)",
    ValueType(): "llvm::PointerType::get(context, 0)",
    CountVarargsValueType(): "llvm::Type::getInt64Ty(context)"
}

lwir(
    template_path = "jitir_llvmapi.tmpl.hpp",
    output_path = "jitir_llvmapi.hpp",
    ir = jitir,
    plugins = [
        LLVMAPIPlugin(
            prefix = "jitir",
            type_substitutions = llvm_type_substitutions
        ),
        BuildBuildInstPlugin(
            type_substitutions = llvm_type_substitutions
        ),
        MapSymbolsInstPlugin(
            prefix = "jitir"
        )
    ],
    placeholder = lambda name: "/* ${" + name + "} */"
)

class MarkdownPlugin:
    def run(self, ir):
        code = ""
        for inst in ir.insts:
            code += f"### {inst.name}\n\n"
            if inst.doc:
                code += f"{inst.doc}\n\n"
            code += f"Arguments\n\n"
            for arg in inst.args:
                if isinstance(arg.type, CountVarargsValueType):
                    code += f"- **{arg.name}**: Variable number of arguments\n"
                else:
                    code += f"- **{arg.name}**: `{arg.type.format(ir)}`\n"
            code += f"\n"
            code += f"Return Type: `{inst.type}`\n\n"
            if len(inst.type_checks) > 0:
                code += f"Type Checks:\n\n"
                for check in inst.type_checks:
                    code += f"- `{check}`\n"
            code += f"\n"
        return {"markdown": code}
        
lwir(
    template_path = "doc/jitir.md.tmpl",
    output_path = "doc/jitir.md",
    ir = jitir,
    plugins = [
        MarkdownPlugin()
    ]
)