generate_types.vsh

import json
import strings

/*
This is a script file which creates static type declarations
for tree-sitter-v's node types by using the information found
in node-types.json and turn it into a pseudo-sum type using enums
and create a NodeTypeFactory implementation that will convert type names
into respective NodeType enum. Anonymous nodes are automatically identified
as NodeType.unknown.

See: https://tree-sitter.github.io/tree-sitter/using-parsers#static-node-types
*/

const to_be_escaped = ['none', 'true', 'false', 'map', 'type']

fn escape_name(name string) string {
	if name in to_be_escaped {
		return name + '_'
	}
	return name
}

fn write_enum_member(mut wr strings.Builder, type_name string, member_name string) {
	wr.write_string('${type_name}.${escape_name(member_name)}')
}

fn write_enum_array(mut wr strings.Builder, enum_type_name string, list []string) {
	wr.writeln('[')
	for i, name in list {
		wr.write_u8(`\t`)
		write_enum_member(mut wr, if i == 0 { enum_type_name } else { '' }, name)
		if i < list.len - 1 {
			wr.write_u8(`,`)
		}
		wr.write_u8(`\n`)
	}
	wr.write_u8(`]`)
}

fn write_const_enum_array(mut wr strings.Builder, var_name string, enum_type_name string, list []string) {
	wr.write_string('\nconst $var_name = ')
	write_enum_array(mut wr, enum_type_name, list)
	wr.write_u8(`\n`)
}

struct NodeType {
	name     string     [json: 'type']
	named    bool
	subtypes []NodeType
	// children TypeChildren
}

fn (ntype NodeType) is_anon() bool {
	return !ntype.named || ntype.name.len == 0 || ntype.name[0] == `_`
}

cur_dir := dir(@FILE)
node_types_json := read_file(join_path(cur_dir, 'src', 'node-types.json'))?
node_types := json.decode([]NodeType, node_types_json)?
node_type_enum_name := 'NodeType'
super_type_enum_name := 'SuperType'
mut file := open_file(join_path(cur_dir, 'node_types.v'), 'w+')?
mut wr := strings.new_builder(1024 * 1024)
mut supertype_node_groups := map[string][]string{}

wr.writeln('module tree_sitter_v\nimport arrays { merge }\n')
wr.writeln('/* This is an auto-generated file. Do not edit this file directly! See `generate_types.vsh` */')

// write supertypes
wr.writeln('pub enum $super_type_enum_name {')
wr.writeln('\tunknown')
for ntype in node_types {
	if !ntype.named || ntype.name.len == 0 || ntype.name[0] != `_` || ntype.subtypes.len == 0 {
		continue
	}
	wr.writeln('\t${escape_name(ntype.name[1..])}')
	supertype_node_groups[ntype.name] = ntype.subtypes.map(it.name)
}
wr.writeln('}\n')

wr.writeln('pub enum $node_type_enum_name {')
wr.writeln('\tunknown\n\terror')

mut declaration_node_types := []string{cap: 100}
mut identifier_node_types := []string{cap: 100}
mut literal_node_types := []string{cap: 100}

// write node types as enum members
for ntype in node_types {
	if ntype.is_anon() {
		continue
	}

	if ntype.name.ends_with('_declaration') {
		declaration_node_types << ntype.name
	} else if ntype.name == 'identifier' || ntype.name.ends_with('_identifier') {
		identifier_node_types << ntype.name
	} else if ntype.name.ends_with('_literal') {
		literal_node_types << ntype.name
	}

	wr.writeln('\t${escape_name(ntype.name)}')
}
wr.writeln('}')

for supertype_name, supertype_node_types in supertype_node_groups {
	wr.write_string('\nconst supertype_${supertype_name}_nodes = ')
	super_type_members := supertype_node_types.filter(it.starts_with('_'))
	for ntype in super_type_members {
		wr.write_string('merge(supertype_${ntype}_nodes, ')
	}
	write_enum_array(mut wr, node_type_enum_name, supertype_node_types.filter(!it.starts_with('_')))
	wr.writeln(')'.repeat(super_type_members.len))
}

wr.write_string('\npub fn (typ $node_type_enum_name) group() $super_type_enum_name {\n\treturn ')
mut if_i := 0
supertype_ordered_names := ['top_level_declaration', 'type', 'simple_type', 'expression',
	'expression_with_blocks', 'statement', 'simple_statement', 'unknown']
for supertype_name in supertype_ordered_names {
	if if_i < supertype_ordered_names.len - 1 {
		wr.write_string('if typ in supertype__${supertype_name}_nodes ')
	}
	wr.write_string('{\n\t\t')
	write_enum_member(mut wr, super_type_enum_name, supertype_name)
	wr.write_string('\n\t}')
	if if_i < supertype_ordered_names.len - 1 {
		wr.write_string(' else ')
	} else {
		wr.write_u8(`\n`)
	}
	if_i++
}
wr.writeln('}')

// write consts
write_const_enum_array(mut wr, 'declaration_node_types', node_type_enum_name, declaration_node_types)
write_const_enum_array(mut wr, 'identifier_node_types', node_type_enum_name, identifier_node_types)
write_const_enum_array(mut wr, 'literal_node_types', node_type_enum_name, literal_node_types)

wr.writeln('\npub fn (typ $node_type_enum_name) is_declaration() bool { return typ in declaration_node_types }')
wr.writeln('pub fn (typ $node_type_enum_name) is_identifier() bool { return typ in identifier_node_types }')
wr.writeln('pub fn (typ $node_type_enum_name) is_literal() bool { return typ in literal_node_types }')

// create VNodeTypeFactory
node_type_factory_sym_name := 'VNodeTypeFactory'

wr.writeln('\npub const type_factory = &$node_type_factory_sym_name{}')
wr.writeln('\npub struct $node_type_factory_sym_name {}')
wr.writeln('\npub fn (nf $node_type_factory_sym_name) get_type(type_name string) $node_type_enum_name {')
wr.writeln('\treturn match type_name {')
wr.writeln("\t\t'ERROR' { ${node_type_enum_name}.error }")

for ntype in node_types {
	if ntype.is_anon() {
		continue
	}
	wr.write_string("\t\t'$ntype.name' { ")
	write_enum_member(mut wr, node_type_enum_name, ntype.name)
	wr.writeln(' }')
}
wr.writeln('\t\t else { ${node_type_enum_name}.unknown }')
wr.writeln('\t}')
wr.writeln('}')

file.write(wr)?
file.close()