package ast

import "../source"
import "../symbol"
import "../types"
import "core:mem"

Expr_Id :: distinct u32
Stmt_Id :: distinct u32
Function_Id :: distinct u32
Global_Id :: distinct u32
Import_Id :: distinct u32
File_Id :: distinct u32
Package_Id :: distinct u32

INVALID_EXPR :: Expr_Id(0xffff_ffff)
INVALID_STMT :: Stmt_Id(0xffff_ffff)
INVALID_FUNCTION :: Function_Id(0xffff_ffff)
INVALID_GLOBAL :: Global_Id(0xffff_ffff)
INVALID_IMPORT :: Import_Id(0xffff_ffff)
INVALID_FILE :: File_Id(0xffff_ffff)
INVALID_PACKAGE :: Package_Id(0xffff_ffff)

expr_id :: proc(index: int) -> Expr_Id {
	assert(index >= 0 && u64(index) < u64(INVALID_EXPR))
	return Expr_Id(index)
}

stmt_id :: proc(index: int) -> Stmt_Id {
	assert(index >= 0 && u64(index) < u64(INVALID_STMT))
	return Stmt_Id(index)
}

function_id :: proc(index: int) -> Function_Id {
	assert(index >= 0 && u64(index) < u64(INVALID_FUNCTION))
	return Function_Id(index)
}

global_id :: proc(index: int) -> Global_Id {
	assert(index >= 0 && u64(index) < u64(INVALID_GLOBAL))
	return Global_Id(index)
}

import_id :: proc(index: int) -> Import_Id {
	assert(index >= 0 && u64(index) < u64(INVALID_IMPORT))
	return Import_Id(index)
}

file_id :: proc(index: int) -> File_Id {
	assert(index >= 0 && u64(index) < u64(INVALID_FILE))
	return File_Id(index)
}

package_id :: proc(index: int) -> Package_Id {
	assert(index >= 0 && u64(index) < u64(INVALID_PACKAGE))
	return Package_Id(index)
}

index :: proc(id: $T, invalid: T, count: int) -> (int, bool) {
	value := int(id)
	return value, id != invalid && value < count
}

Type_Syntax :: types.Type

Type_Use :: struct {
	type: Type_Syntax,
	span: source.Span,
	pkg:  Package_Id,
	file: File_Id,
	diagnostic: source.Diagnostic_Id,
}

Expr_Kind :: enum u8 {
	Invalid,
	Integer,
	Float,
	String,
	Bool,
	Array,
	Null,
	Undefined,
	Inference_Hole,
	Type,
	Name,
	Enum_Literal,
	Address,
	Deref,
	Index,
	Slice,
	Field,
	Unwrap,
	Orelse,
	Struct_Literal,
	Keyed,
	Cast,
	Comptime,
	Negate,
	Not,
	Bit_Not,
	Add,
	Sub,
	Mul,
	Div,
	Bit_And,
	Bit_Or,
	Bit_Xor,
	Shift_Left,
	Shift_Right,
	Shift_Left_Saturating,
	Eq,
	Ne,
	Lt,
	Le,
	Gt,
	Ge,
	And,
	Or,
	Range,
	Call,
	Try,
	Catch,
	Function_Literal,
	Anonymous_Struct_Type,
}

Expr :: struct {
	span:       source.Span,
	integer:    u64,
	args:       []Expr_Id,
	qualifier:  symbol.Id,
	name:       symbol.Id,
	type:       Type_Syntax,
	left:       Expr_Id,
	right:      Expr_Id,
	body:       []Stmt_Id,
	diagnostic: source.Diagnostic_Id,
	parenthesized: bool,
	intrinsic:   bool,
	tuple:       bool,
	kind:       Expr_Kind,
}

Param :: struct {
	name: symbol.Id,
	span: source.Span,
	type: Type_Syntax,
	comptime_value: bool,
}

Stmt_Kind :: enum u8 {
	Invalid,
	Declaration,
	Assignment,
	Return,
	Expression,
	If,
	While,
	For,
	Break,
	Continue,
	Block,
	Defer,
	Yield,
	Match,
	Match_Arm,
}

Assignment_Op :: enum u8 {
	Set,
	Add,
	Sub,
	Mul,
	Div,
	Bit_And,
	Bit_Or,
	Bit_Xor,
	Shift_Left,
	Shift_Right,
	Shift_Left_Saturating,
}

Stmt :: struct {
	kind:       Stmt_Kind,
	span:       source.Span,
	name:       symbol.Id,
	index_name: symbol.Id,
	// `For`/`While` loop bodies may carry a label (`blk: { ... }`); a `Yield`
	// may target one (`yield :blk x`), letting a yield reach past an enclosing
	// `if` to exit the labeled loop. `INVALID` when absent.
	label:      symbol.Id,
	type:       Type_Syntax,
	immutable:  bool,
	value_control_flow: bool,
	pointer_capture: bool,
	expand: bool,
	error_only: bool,
	// Assignments store the lvalue in `target`, the right-hand side in `expr`,
	// and the source operator in `assignment_op`. `Set` is ordinary `=`;
	// the remaining variants preserve their corresponding compound operator.
	assignment_op: Assignment_Op,
	target:     Expr_Id,
	expr:       Expr_Id,
	// `If` statements use `expr` as the condition, `captures` as optional
	// unwrap binding names, `guard` as the optional post-unwrap boolean
	// condition, `body` as the then-block, and `else_body` as the else-block.
	// An `else if` chain is represented as an `else_body` holding a single
	// nested `If` statement.
	// `While` statements use `expr` as the condition, `body` as the loop body,
	// and `update` as the optional post-iteration statement.
	// `For` statements use `expr` as the iterable, `name` as the item capture,
	// `index_name` as the optional index capture, and `pointer_capture` to
	// distinguish `|@item|` from copy capture.
	// `Block` statements (a bare `{ ... }` scope) use `body` as their statements.
	// `Defer` statements use `update` as the deferred statement (which may itself
	// be a `Block`), `error_only` for `errdefer`, and `captures` for its optional
	// error capture.
	// `Match` statements use `expr` as the subject and `body` as the list of arm
	// statements (each a `Match_Arm`). A `Match_Arm` uses `patterns` as its pattern
	// list (empty marks the `else` arm; more than one is a multi-pattern arm),
	// `captures` for the optional payload capture (0 or 1 name, tagged-union variants
	// only) with `pointer_capture` distinguishing `|@cap|` from `|cap|`, and `body`
	// as the arm body. An expanded arm has `expand` set, no patterns, and one or two
	// captures for its specialized value/payload and optional tagged-union tag.
	captures:   []symbol.Id,
	// `Match_Arm` pattern list; empty ⇒ the `else` arm.
	patterns:   []Expr_Id,
	guard:      Expr_Id,
	body:       []Stmt_Id,
	else_body:  []Stmt_Id,
	update:     Stmt_Id,
	diagnostic: source.Diagnostic_Id,
}

Function :: struct {
	span:       source.Span,
	name:       symbol.Id,
	pkg:        Package_Id,
	file:       File_Id,
	c_abi:      bool,
	imported:   bool,
	generated:  bool,
	analysis_root: bool,
	test:       bool,
	file_hidden: bool,
	has_body:   bool,
	variadic:   bool,
	params:     []Param,
	result:     Type_Syntax,
	error:      Type_Syntax,
	body:       []Stmt_Id,
	link_name:  string,
	unsupported_reason: string,
	diagnostic: source.Diagnostic_Id,
}

Global :: struct {
	span:       source.Span,
	name:       symbol.Id,
	link_name:  string,
	pkg:        Package_Id,
	file:       File_Id,
	type:       Type_Syntax,
	immutable:  bool,
	external:   bool,
	writable:   bool,
	file_hidden: bool,
	expr:       Expr_Id,
	diagnostic: source.Diagnostic_Id,
}

Enum_Value :: struct {
	expr:     Expr_Id,
	span:     source.Span,
	explicit: bool,
}

Enum_Declaration :: struct {
	type:   types.Type,
	pkg:    Package_Id,
	file:   File_Id,
	values: []Enum_Value,
}

Import :: struct {
	span:       source.Span,
	alias:      symbol.Id,
	path:       string,
	pkg:        Package_Id,
	file:       File_Id,
	target:     Package_Id,
	valid:      bool,
	used:       bool,
	test_only:  bool,
	diagnostic: source.Diagnostic_Id,
}

Declaration_Alias_Kind :: enum u8 {
	Invalid,
	Function,
	Global,
	Type,
}

Declaration_Alias :: struct {
	span:        source.Span,
	name:        symbol.Id,
	qualifier:   symbol.Id,
	member:      symbol.Id,
	pkg:         Package_Id,
	file:        File_Id,
	target_pkg:  Package_Id,
	target:      u32,
	kind:        Declaration_Alias_Kind,
	file_hidden: bool,
	valid:       bool,
	diagnostic:  source.Diagnostic_Id,
}

Struct_Field_Default :: struct {
	record: Type_Syntax,
	field:  symbol.Id,
	expr:   Expr_Id,
	pkg:    Package_Id,
	file:   File_Id,
}

File :: struct {
	source:  source.Source_Id,
	pkg:     Package_Id,
}

Package :: struct {
	path:      string,
	name:      symbol.Id,
	available: bool,
	test:      bool,
	kind:      Package_Kind,
}

Compile_Mode :: enum u8 {
	Executable,
	Test,
}

Package_Kind :: enum u8 {
	Native,
	C_Header,
}

Unsupported :: struct {
	pkg:    Package_Id,
	name:   symbol.Id,
	reason: string,
}

// Trampoline is a generated C wrapper source that must be compiled and linked
// alongside the program so an internal-linkage (`static inline`) C function can
// be called through an external symbol. `source` is the wrapper function only;
// `header` is the absolute path it must `#include` (emitted once per header).
Trampoline :: struct {
	symbol: string,
	source: string,
	header: string,
}

Module :: struct {
	exprs:      [dynamic]Expr,
	statements: [dynamic]Stmt,
	functions:  [dynamic]Function,
	globals:    [dynamic]Global,
	enum_declarations: [dynamic]Enum_Declaration,
	imports:    [dynamic]Import,
	aliases:    [dynamic]Declaration_Alias,
	files:      [dynamic]File,
	packages:   [dynamic]Package,
	unsupported: [dynamic]Unsupported,
	c_trampolines: [dynamic]Trampoline,
	strings:    [dynamic]string,
	type_fields: [dynamic]types.Field,
	struct_field_defaults: [dynamic]Struct_Field_Default,
	type_uses:   [dynamic]Type_Use,
	type_store: types.Store,
	allocator:  mem.Allocator,
}

init_module :: proc(allocator := context.allocator) -> Module {
	module: Module
	module.allocator = allocator
	module.type_store = types.init_store(allocator)
	module.exprs.allocator = allocator
	module.statements.allocator = allocator
	module.functions.allocator = allocator
	module.globals.allocator = allocator
	module.enum_declarations.allocator = allocator
	module.imports.allocator = allocator
	module.aliases.allocator = allocator
	module.files.allocator = allocator
	module.packages.allocator = allocator
	module.unsupported.allocator = allocator
	module.c_trampolines.allocator = allocator
	module.strings.allocator = allocator
	module.type_fields.allocator = allocator
	module.struct_field_defaults.allocator = allocator
	module.type_uses.allocator = allocator
	return module
}

destroy_module :: proc(module: ^Module) {
	for expr in module.exprs {
		delete(expr.args, module.allocator)
		delete(expr.body, module.allocator)
	}
	for statement in module.statements {
		delete(statement.captures, module.allocator)
		delete(statement.patterns, module.allocator)
		delete(statement.body, module.allocator)
		delete(statement.else_body, module.allocator)
	}
	for function in module.functions {
		delete(function.params, module.allocator)
		delete(function.body, module.allocator)
		delete(function.link_name, module.allocator)
		delete(function.unsupported_reason, module.allocator)
	}
	for import_item in module.imports {
		delete(import_item.path, module.allocator)
	}
	for global in module.globals {
		delete(global.link_name, module.allocator)
	}
	for declaration in module.enum_declarations {
		delete(declaration.values, module.allocator)
	}
	for pkg in module.packages {
		delete(pkg.path, module.allocator)
	}
	for item in module.unsupported {
		delete(item.reason, module.allocator)
	}
	for trampoline in module.c_trampolines {
		delete(trampoline.symbol, module.allocator)
		delete(trampoline.source, module.allocator)
		delete(trampoline.header, module.allocator)
	}
	for value in module.strings {
		delete(value, module.allocator)
	}
	delete(module.exprs)
	delete(module.statements)
	delete(module.functions)
	delete(module.globals)
	delete(module.enum_declarations)
	delete(module.imports)
	delete(module.aliases)
	delete(module.files)
	delete(module.packages)
	delete(module.unsupported)
	delete(module.c_trampolines)
	delete(module.strings)
	delete(module.type_fields)
	delete(module.struct_field_defaults)
	delete(module.type_uses)
	types.destroy_store(&module.type_store)
}