Ink-Shader-Language/AST.jai

/////////////////////////////////////
//~ nbr: Node data structure
//
// [ ] Add a way to infer or get file path from a node

AST_Kind :: enum {
	Program;
	Function;
	Return;

	// @Incomplete(nb): Should these three really be their own block types?
	//                  Maybe they at least shouldn't need to have their own tokens...
	Properties;
	Meta;
	Instance;
	//==

	// Hint;
	// Type;
	// Operator;
	Call;
	Struct;
	If;
	CBuffer;
	FieldList;
	ArgList;
	Variable;
	Binary;
	Unary;
	Integer;
	Float;
	Expression_Statement;
	Field;
	Unnamed_Field;
	Block;
	Error;
}

AST_Node :: struct {
	kind : AST_Kind;

	// @Note(niels): Children nodes can be interpreted as anything useful.
	// for an if-statement we would have at most 2 children
	// a property block has a field list child node which has
	// a child node for each field declaration etc.
	children : [..]*AST_Node;
	parent : *AST_Node;

	// @Note(niels): Every node can have a name, but most nodes don't. A function or field declaration has one,
	// but an if-statement does not
	name     : string;

	integer_value  : int;
	float_value    : float;

	token : Token;

	array_field : bool;

	source_location : Source_Range;

	type_variable : Type_Variable_Handle;

	foreign_declaration : bool;

	// @Incomplete(nb): Change this to just be children and a single token_data field,
	//                  then we can add new node types (hint, typespec, operator) and have them
	//                  as children instead.
	hint_tokens    : [..]Token;
	
	vertex_entry_point : bool;
	pixel_entry_point  : bool;
}


// ===========================================================
// Pretty printing
pretty_print_call :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
	indent(builder, indentation);
	append(builder, "(");
	append(builder, node.name);
	if node.children.count > 0 {
		append(builder, " ");
		pretty_print_children(node.children[0], indentation, builder, flags = 0);
	}
	append(builder, ")");
}

pretty_print_arglist :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
	indent(builder, indentation);
	append(builder, "[");

	pretty_print_children(node, indentation + 1, builder, flags = .NewLine);

	append(builder, "]");
}

pretty_print_fieldlist :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
	indent(builder, indentation);

	append(builder, "[");
	pretty_print_children(node, indentation + 1, builder, flags = .NewLine);

	append(builder, "]");
}

pretty_print_field :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
	print_to_builder(builder, tprint("(:= %", node.name));

	if node.kind != .Unnamed_Field && node.token.ident_value.count > 0 {
		if node.array_field {
			append(builder, " [");
			pretty_print_node(node.children[0], 0, builder);
			append(builder, "].");
			print_to_builder(builder, "%", node.token.ident_value);
		} else {
			print_to_builder(builder, " %", node.token.ident_value);
		}
		
	}

	for hint : node.hint_tokens {
		if hint.string_value.count > 0 {
			print_to_builder(builder, " (@%)", hint.string_value);
		}
	}

	if !node.array_field && node.children.count > 0 {
		append(builder, " ");
		pretty_print_children(node, indentation, builder);
	}

	append(builder, ")");
}

Children_Print_Flags :: enum_flags {
	NewLine   :: 1 << 0;
	Separator :: 1 << 1;
	Space     :: 1 << 2;
}

pretty_print_children :: (parent : *AST_Node, indentation : int, builder : *String_Builder, flags : Children_Print_Flags = .Separator) {
	if !parent {
		return;
	}
	
	children := parent.children;
	for child : children {
		if it_index > 0 {
			indent(builder, indentation);
		}

		if !child continue;
		pretty_print_node(child, 0, builder);
		
		if it_index != children.count - 1 {
			if flags & .Separator {
				append(builder, ",");
			}

			append(builder, " ");

			if flags & .NewLine {
				append(builder, "\n");
			}
		}
	}
}

op_to_string :: (oper : Token) -> string {
	if oper.kind == {
		case .TOKEN_PLUS;
		return "+";
		case .TOKEN_MINUS;
		return "-";
		case .TOKEN_STAR;
		return "*";
		case .TOKEN_SLASH;
		return "/";
		case .TOKEN_ISEQUAL;
		return "==";
		case .TOKEN_ASSIGN;
		return "=";
		case .TOKEN_ISNOTEQUAL;
		return "!=";
		case .TOKEN_LOGICALOR;
		return "||";
		case .TOKEN_LOGICALAND;
		return "&&";
		case .TOKEN_LESS;
		return "<";
		case .TOKEN_LESSEQUALS;
		return "<=";
		case .TOKEN_GREATER;
		return ">";
		case .TOKEN_GREATEREQUALS;
		return ">=";
	}
	return "";
}

pretty_print_binary :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
	indent(builder, indentation);
	append(builder, "(");
	op := node.token;

	print_to_builder(builder, op_to_string(op));
	append(builder, " ");

	pretty_print_children(node, 0, builder, flags = 0);
	append(builder, ")");
}

pretty_print_unary :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
	indent(builder, indentation);
	op := node.token;

	print_to_builder(builder, op_to_string(op));
	pretty_print_children(node, 0, builder, flags = 0);
}

print_return_node :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
	append(builder, "(return ");

	pretty_print_children(node, 0, builder);

	append(builder, ")");
}

pretty_print_if :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
indent(builder, indentation);
	append(builder, "(if ");

	body := node.children[0];
	pretty_print_node(body, 0, builder);

	if node.children.count == 2 {
		append(builder, "\n");
		indent(builder, indentation);
		append(builder, "(else ");
		pretty_print_node(node.children[1], 0, builder);
		append(builder, ")");
	}

	append(builder, ")");
}

print_expression_statement :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
	indent(builder, indentation);

	if node.children[0] {
		pretty_print_node(node.children[0], indentation, builder);
	}
}

pretty_print_node :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
	if node.kind == {
		case .Return; {
			print_return_node(node, indentation, builder);
		}
		case .If; {
			pretty_print_if(node, indentation, builder);
		}
		case .Struct;
		case .ArgList; {
			pretty_print_arglist(node, indentation + 2, builder);
		}
		case .FieldList; {
			pretty_print_fieldlist(node, indentation + 2, builder);
		}
		case .Field; {
			pretty_print_field(node, indentation, builder);
		}
		case .Unnamed_Field; {
			pretty_print_field(node, indentation, builder);
		}
		case .Block; {
			pretty_print_children(node, indentation + 2, builder, flags = .NewLine);
		}
		case .Binary; {
			pretty_print_binary(node, indentation, builder);
		}
		case .Unary; {
			pretty_print_unary(node, indentation, builder);
		}
		case .Variable; {
			pretty_print_variable(node, indentation, builder);
		}
		case .Expression_Statement; {
			print_expression_statement(node, indentation, builder);
		}
		case .Integer; {
			print_to_builder(builder, "%", node.integer_value);
		}
		case .Float; {
			print_to_builder(builder, "%", node.float_value);
		}
		case .Call; {
			pretty_print_call(node, indentation, builder);
		}
		case .Error; {
			print_to_builder(builder, "(error \"%\")", node.name);
		}
	}
}

pretty_print_variable :: (node : *AST_Node, indentation : int, builder : *String_Builder) {
	indent(builder, indentation);
	print_to_builder(builder, "%", node.name);
	
	for child : node.children {
		if child.kind == .Variable {
			append(builder, ".");
			pretty_print_variable(child, indentation, builder);
		} else if child.kind == .Unary {
			append(builder, "[");
			pretty_print_node(child.children[0], 0, builder);
			append(builder, "]");
		}
	}
}

pretty_print_declaration :: (declaration : *AST_Node, indentation : int, builder : *String_Builder) {
	indent(builder, indentation);
	append(builder, "(");

	if declaration.foreign_declaration {
		append(builder, "foreign ");
	}

	if declaration.kind == .Function {
		append(builder, "fun ");
	}
	
	if declaration.vertex_entry_point {
		append(builder, "vertex ");
	}

	if declaration.pixel_entry_point {
		append(builder, "pixel ");
	}

	if declaration.kind == .Properties {
		append(builder, "properties");
		if declaration.name.count > 0 {
			print_to_builder(builder, " %", declaration.name);
		}
	} else if declaration.kind == .Instance {
		append(builder, "instance");
	} else if declaration.kind == .Meta {
		append(builder, "meta");
	}
	else {
		if declaration.kind == .Struct {
			append(builder, "struct ");
		} else if declaration.kind == .CBuffer {
			append(builder, "constant_buffer ");
		}
		print_to_builder(builder, "%", declaration.name);
	}

	if declaration.kind == .Function && declaration.token.kind == .TOKEN_IDENTIFIER{
		print_to_builder(builder, " -> %", declaration.token.ident_value);
		for hint : declaration.hint_tokens {
			if hint.string_value.count > 0 {
				print_to_builder(builder, " (@%)", hint.string_value);
			}
		}

	}

	if declaration.children.count > 0 {
		print_to_builder(builder, "\n");
		pretty_print_children(declaration, indentation + 1, builder, flags = .NewLine);
	}

	append(builder, ")");

}

pretty_print_ast :: (root : *AST_Node, allocator : Allocator) -> string {
	builder : String_Builder;
	init_string_builder(*builder,, allocator);

	indentation := 0;
	
	append(*builder, "(");
	append(*builder, "program\t\n");

	indentation += 1;

	declarations := root.children;

	for declaration : declarations {
		pretty_print_declaration(declaration, indentation, *builder);

		if it_index < declarations.count - 1 {
			append(*builder, "\n\n");
		}
	}

	append(*builder, ")");

	return builder_to_string(*builder,, allocator);
}

#scope_file
indent :: (builder : *String_Builder, indentation : int) {
	for 0..indentation - 1 {
		append(builder, " ");
	}
}