package odin_ast

import "core:fmt"

// A Visitor's visit procedure is invoked for each node encountered by walk
// If the result visitor is not nil, walk visits each of the children of node with the new visitor,
// followed by a call of v.visit(v, nil)
Visitor :: struct {
	visit: proc(visitor: ^Visitor, node: ^Node) -> ^Visitor,
	data:  rawptr,
}


// inspect traverses an AST in depth-first order
// It starts by calling f(node), and node must be non-nil
// If f returns true, inspect invokes f recursively for each of the non-nil children of node,
// followed by a call of f(nil)
inspect :: proc(node: ^Node, f: proc(^Node) -> bool) {
	v := &Visitor{
		visit = proc(v: ^Visitor, node: ^Node) -> ^Visitor {
			f := (proc(^Node) -> bool)(v.data)
			if f(node) {
				return v
			}
			return nil
		},
		data = rawptr(f),
	}
	walk(v, node)
}



// walk traverses an AST in depth-first order: It starts by calling v.visit(v, node), and node must not be nil
// If the visitor returned by v.visit(v, node) is not nil, walk is invoked recursively with the new visitor for
// each of the non-nil children of node, followed by a call of the new visit procedure
walk :: proc(v: ^Visitor, node: ^Node) {
	walk_expr_list :: proc(v: ^Visitor, list: []^Expr) {
		for x in list {
			walk(v, x)
		}
	}

	walk_stmt_list :: proc(v: ^Visitor, list: []^Stmt) {
		for x in list {
			walk(v, x)
		}
	}
	walk_attribute_list :: proc(v: ^Visitor, list: []^Attribute) {
		for x in list {
			walk(v, x)
		}
	}


	v := v
	if v = v->visit(node); v == nil {
		return
	}

	switch n in &node.derived {
	case File:
		if n.docs != nil {
			walk(v, n.docs)
		}
		walk_stmt_list(v, n.decls[:])
	case Package:
		for _, f in n.files {
			walk(v, f)
		}

	case Comment_Group:
		// empty
	case Bad_Expr:
	case Ident:
	case Implicit:
	case Undef:
	case Basic_Lit:
	case Basic_Directive:
	case Ellipsis:
		if n.expr != nil {
			walk(v, n.expr)
		}
	case Proc_Lit:
		walk(v, n.type)
		walk(v, n.body)
		walk_expr_list(v, n.where_clauses)
	case Comp_Lit:
		if n.type != nil {
			walk(v, n.type)
		}
		walk_expr_list(v, n.elems)
	case Tag_Expr:
		walk(v, n.expr)
	case Unary_Expr:
		walk(v, n.expr)
	case Binary_Expr:
		walk(v, n.left)
		walk(v, n.right)
	case Paren_Expr:
		walk(v, n.expr)
	case Selector_Expr:
		walk(v, n.expr)
		walk(v, n.field)
	case Implicit_Selector_Expr:
		walk(v, n.field)
	case Selector_Call_Expr:
		walk(v, n.expr)
		walk(v, n.call)
	case Index_Expr:
		walk(v, n.expr)
		walk(v, n.index)
	case Matrix_Index_Expr:
		walk(v, n.expr)
		walk(v, n.row_index)
		walk(v, n.column_index)
	case Deref_Expr:
		walk(v, n.expr)
	case Slice_Expr:
		walk(v, n.expr)
		if n.low != nil {
			walk(v, n.low)
		}
		if n.high != nil {
			walk(v, n.high)
		}
	case Call_Expr:
		walk(v, n.expr)
		walk_expr_list(v, n.args)
	case Field_Value:
		walk(v, n.field)
		walk(v, n.value)
	case Ternary_If_Expr:
		walk(v, n.x)
		walk(v, n.cond)
		walk(v, n.y)
	case Ternary_When_Expr:
		walk(v, n.x)
		walk(v, n.cond)
		walk(v, n.y)
	case Or_Else_Expr:
		walk(v, n.x)
		walk(v, n.y)
	case Or_Return_Expr:
		walk(v, n.expr)
	case Type_Assertion:
		walk(v, n.expr)
		if n.type != nil {
			walk(v, n.type)
		}
	case Type_Cast:
		walk(v, n.type)
		walk(v, n.expr)
	case Auto_Cast:
		walk(v, n.expr)
	case Inline_Asm_Expr:
		walk_expr_list(v, n.param_types)
		walk(v, n.return_type)
		walk(v, n.constraints_string)
		walk(v, n.asm_string)


	case Bad_Stmt:
	case Empty_Stmt:
	case Expr_Stmt:
		walk(v, n.expr)
	case Tag_Stmt:
		walk(v, n.stmt)
	case Assign_Stmt:
		walk_expr_list(v, n.lhs)
		walk_expr_list(v, n.rhs)
	case Block_Stmt:
		if n.label != nil {
			walk(v, n.label)
		}
		walk_stmt_list(v, n.stmts)
	case If_Stmt:
		if n.label != nil {
			walk(v, n.label)
		}
		if n.init != nil {
			walk(v, n.init)
		}
		walk(v, n.cond)
		walk(v, n.body)
		if n.else_stmt != nil {
			walk(v, n.else_stmt)
		}
	case When_Stmt:
		walk(v, n.cond)
		walk(v, n.body)
		if n.else_stmt != nil {
			walk(v, n.else_stmt)
		}
	case Return_Stmt:
		walk_expr_list(v, n.results)
	case Defer_Stmt:
		walk(v, n.stmt)
	case For_Stmt:
		if n.label != nil {
			walk(v, n.label)
		}
		if n.init != nil {
			walk(v, n.init)
		}
		if n.cond != nil {
			walk(v, n.cond)
		}
		if n.post != nil {
			walk(v, n.post)
		}
		walk(v, n.body)
	case Range_Stmt:
		if n.label != nil {
			walk(v, n.label)
		}
		for val in n.vals {
			if val != nil {
				walk(v, val)
			}
		}
		walk(v, n.expr)
		walk(v, n.body)
	case Inline_Range_Stmt:
		if n.label != nil {
			walk(v, n.label)
		}
		if n.val0 != nil {
			walk(v, n.val0)
		}
		if n.val1 != nil {
			walk(v, n.val1)
		}
		walk(v, n.expr)
		walk(v, n.body)
	case Case_Clause:
		walk_expr_list(v, n.list)
		walk_stmt_list(v, n.body)
	case Switch_Stmt:
		if n.label != nil {
			walk(v, n.label)
		}
		if n.init != nil {
			walk(v, n.init)
		}
		if n.cond != nil {
			walk(v, n.cond)
		}
		walk(v, n.body)
	case Type_Switch_Stmt:
		if n.label != nil {
			walk(v, n.label)
		}
		if n.tag != nil {
			walk(v, n.tag)
		}
		if n.expr != nil {
			walk(v, n.expr)
		}
		walk(v, n.body)
	case Branch_Stmt:
		if n.label != nil {
			walk(v, n.label)
		}
	case Using_Stmt:
		walk_expr_list(v, n.list)


	case Bad_Decl:
	case Value_Decl:
		if n.docs != nil {
			walk(v, n.docs)
		}
		walk_attribute_list(v, n.attributes[:])
		walk_expr_list(v, n.names)
		if n.type != nil {
			walk(v, n.type)
		}
		walk_expr_list(v, n.values)
		if n.comment != nil {
			walk(v, n.comment)
		}
	case Package_Decl:
		if n.docs != nil {
			walk(v, n.docs)
		}
		if n.comment != nil {
			walk(v, n.comment)
		}
	case Import_Decl:
		if n.docs != nil {
			walk(v, n.docs)
		}
		if n.comment != nil {
			walk(v, n.comment)
		}
	case Foreign_Block_Decl:
		if n.docs != nil {
			walk(v, n.docs)
		}
		walk_attribute_list(v, n.attributes[:])
		if n.foreign_library != nil {
			walk(v, n.foreign_library)
		}
		walk(v, n.body)
	case Foreign_Import_Decl:
		if n.docs != nil {
			walk(v, n.docs)
		}
		walk_attribute_list(v, n.attributes[:])
		walk(v, n.name)
		if n.comment != nil {
			walk(v, n.comment)
		}

	case Proc_Group:
		walk_expr_list(v, n.args)
	case Attribute:
		walk_expr_list(v, n.elems)
	case Field:
		if n.docs != nil {
			walk(v, n.docs)
		}
		walk_expr_list(v, n.names)
		if n.type != nil {
			walk(v, n.type)
		}
		if n.default_value != nil {
			walk(v, n.default_value)
		}
		if n.comment != nil {
			walk(v, n.comment)
		}
	case Field_List:
		for x in n.list {
			walk(v, x)
		}
	case Typeid_Type:
		if n.specialization != nil {
			walk(v, n.specialization)
		}
	case Helper_Type:
		walk(v, n.type)
	case Distinct_Type:
		walk(v, n.type)
	case Poly_Type:
		walk(v, n.type)
		if n.specialization != nil {
			walk(v, n.specialization)
		}
	case Proc_Type:
		walk(v, n.params)
		walk(v, n.results)
	case Pointer_Type:
		walk(v, n.elem)
	case Multi_Pointer_Type:
		walk(v, n.elem)
	case Array_Type:
		if n.tag != nil {
			walk(v, n.tag)
		}
		if n.len != nil {
			walk(v, n.len)
		}
		walk(v, n.elem)
	case Dynamic_Array_Type:
		if n.tag != nil {
			walk(v, n.tag)
		}
		walk(v, n.elem)
	case Struct_Type:
		if n.poly_params != nil {
			walk(v, n.poly_params)
		}
		if n.align != nil {
			walk(v, n.align)
		}
		walk_expr_list(v, n.where_clauses)
		walk(v, n.fields)
	case Union_Type:
		if n.poly_params != nil {
			walk(v, n.poly_params)
		}
		if n.align != nil {
			walk(v, n.align)
		}
		walk_expr_list(v, n.where_clauses)
		walk_expr_list(v, n.variants)
	case Enum_Type:
		if n.base_type != nil {
			walk(v, n.base_type)
		}
		walk_expr_list(v, n.fields)
	case Bit_Set_Type:
		walk(v, n.elem)
		if n.underlying != nil {
			walk(v, n.underlying)
		}
	case Map_Type:
		walk(v, n.key)
		walk(v, n.value)
	case Relative_Type:
		walk(v, n.tag)
		walk(v, n.type)
	case Matrix_Type:
		walk(v, n.row_count)
		walk(v, n.column_count)
		walk(v, n.elem)

	case:
		fmt.panicf("ast.walk: unexpected node type %T", n)
	}

	v->visit(nil)
}

