-- Copyright 2026 Artur Ugnivenko. See LICENSE.
-- Odin LPeg lexer.

-- References:
-- https://odin-lang.org/spec/grammar/
-- https://pkg.odin-lang.org/base/builtin/

local lexer = lexer
local P, S, R = lpeg.P, lpeg.S, lpeg.R

local lex = lexer.new(...)

local ws = S(' \t')^0
local word_end = -(lexer.alnum + '_')

-- Keywords.
lex:add_rule('keyword', lex:tag(lexer.KEYWORD, lex:word_match(lexer.KEYWORD)))

-- Definitions: `name :: proc` and `Name :: struct / enum / union / distinct ...`.
local proc_def = lex:tag(lexer.FUNCTION, lexer.word) *
	#(ws * '::' * ws * ('#' * lexer.word * ws)^0 * 'proc' * word_end)
local type_def = lex:tag(lexer.TYPE, lexer.word) *
	#(ws * '::' * ws *
		(P('struct') + 'enum' + 'union' + 'bit_set' + 'bit_field' + 'distinct') * word_end)
lex:add_rule('definition', proc_def + type_def)

-- Types.
lex:add_rule('type', lex:tag(lexer.TYPE, lex:word_match(lexer.TYPE)))

-- Polymorphic parameters: $T, $N.
lex:add_rule('poly', lex:tag(lexer.TYPE, '$' * lexer.word))

-- Functions.
local builtin_func = -lpeg.B('.') *
	lex:tag(lexer.FUNCTION_BUILTIN, lex:word_match(lexer.FUNCTION_BUILTIN))
local func = lex:tag(lexer.FUNCTION, lexer.word)
lex:add_rule('function', (builtin_func + func) * #P('('))

-- Constants.
lex:add_rule('constant', lex:tag(lexer.CONSTANT_BUILTIN, lex:word_match(lexer.CONSTANT_BUILTIN)))

-- Variables.
lex:add_rule('variable',
	lex:tag(lexer.VARIABLE_BUILTIN, lex:word_match(lexer.VARIABLE_BUILTIN)))

-- Identifiers.
lex:add_rule('identifier', lex:tag(lexer.IDENTIFIER, lexer.word))

-- Strings. Quoted strings end at the line; raw strings span lines, no escapes.
local sq_str = lexer.range("'", true)
local dq_str = lexer.range('"', true)
local raw_str = lexer.range('`', false, false)
lex:add_rule('string', lex:tag(lexer.STRING, sq_str + dq_str + raw_str))

-- Comments.
local comment = lexer.to_eol(P('//') + '#!')
local block_comment = lexer.range('/*', '*/', false, false, true)
local file_tag = lexer.to_eol('#+')
lex:add_rule('comment', lex:tag(lexer.COMMENT, comment + block_comment + file_tag))

-- Directives.
lex:add_rule('directive', lex:tag(lexer.PREPROCESSOR, '#' * S(' \t')^0 * lexer.word))

-- Attributes.
lex:add_rule('attribute', lex:tag(lexer.PREPROCESSOR, '@' * lexer.word^-1))

-- Numbers: 1_000_000, 1.5e-3, 0xff, 0b1010, 0o17, 0z1a (dozenal), 0h3f80_0000 (hex float), 2i.
local function digits(d) return d * (P('_')^-1 * d)^0 end
local dec = digits(lexer.digit)
local float = (dec * '.' * dec^-1 * -P('.') + '.' * dec) * (S('eE') * S('+-')^-1 * dec)^-1 +
	dec * S('eE') * S('+-')^-1 * dec
local int = '0' * (
	S('xh') * P('_')^-1 * digits(lexer.xdigit) +
	'b' * P('_')^-1 * digits(S('01')) +
	'o' * P('_')^-1 * digits(R('07')) +
	'z' * P('_')^-1 * digits(R('09') + S('abAB'))) + dec
lex:add_rule('number', lex:tag(lexer.NUMBER, (float + int) * S('ijk')^-1))

-- Operators.
lex:add_rule('operator', lex:tag(lexer.OPERATOR, '..' + S('+-/*%<>!=^&|?~:;,.()[]{}')))

-- Fold points.
lex:add_fold_point(lexer.OPERATOR, '(', ')')
lex:add_fold_point(lexer.OPERATOR, '[', ']')
lex:add_fold_point(lexer.OPERATOR, '{', '}')
lex:add_fold_point(lexer.COMMENT, '/*', '*/')

-- Word lists.
lex:set_word_list(lexer.KEYWORD, {
	'asm', 'auto_cast', 'bit_field', 'bit_set', 'break', 'case', 'cast', 'continue', 'defer',
	'distinct', 'do', 'dynamic', 'else', 'enum', 'fallthrough', 'for', 'foreign', 'if', 'import',
	'in', 'map', 'matrix', 'not_in', 'or_break', 'or_continue', 'or_else', 'or_return', 'package',
	'proc', 'return', 'struct', 'switch', 'transmute', 'union', 'using', 'when', 'where'
})

lex:set_word_list(lexer.TYPE, {
	'any', 'b8', 'b16', 'b32', 'b64', 'bool', 'byte', 'complex32', 'complex64', 'complex128',
	'cstring', 'cstring16', 'f16', 'f16le', 'f16be', 'f32', 'f32le', 'f32be', 'f64', 'f64le', 'f64be',
	'i8', 'i16', 'i16le', 'i16be', 'i32', 'i32le', 'i32be', 'i64', 'i64le', 'i64be', 'i128',
	'i128le', 'i128be', 'int', 'quaternion64', 'quaternion128', 'quaternion256', 'rawptr', 'rune',
	'string', 'string16', 'typeid', 'u8', 'u16', 'u16le', 'u16be', 'u32', 'u32le', 'u32be', 'u64',
	'u64le', 'u64be', 'u128', 'u128le', 'u128be', 'uint', 'uintptr'
})

lex:set_word_list(lexer.FUNCTION_BUILTIN, {
	-- compiler built-ins
	'abs', 'align_of', 'cap', 'clamp', 'complex', 'compress_values', 'conj', 'expand_values', 'imag',
	'jmag', 'kmag', 'len', 'max', 'min', 'offset_of', 'offset_of_by_string', 'offset_of_member',
	'offset_of_selector', 'quaternion', 'raw_data', 'real', 'size_of', 'soa_unzip', 'soa_zip',
	'swizzle', 'type_info_of', 'type_of', 'typeid_of',
	-- base:runtime procs exported as @(builtin)
	'append', 'assert', 'card', 'clear', 'copy', 'delete', 'delete_key', 'ensure', 'free',
	'free_all', 'inject_at', 'make', 'new', 'new_clone', 'ordered_remove', 'panic', 'pop',
	'pop_front', 'reserve', 'resize', 'shrink', 'unimplemented', 'unordered_remove', 'unreachable'
})

lex:set_word_list(lexer.CONSTANT_BUILTIN, {
	'false', 'true', 'nil', 'ODIN_ARCH', 'ODIN_BUILD_MODE', 'ODIN_COMPILE_TIMESTAMP', 'ODIN_DEBUG',
	'ODIN_DEFAULT_TO_NIL_ALLOCATOR', 'ODIN_DEFAULT_TO_PANIC_ALLOCATOR', 'ODIN_DISABLE_ASSERT',
	'ODIN_ENDIAN', 'ODIN_ERROR_POS_STYLE', 'ODIN_NO_BOUNDS_CHECK', 'ODIN_NO_CRT',
	'ODIN_NO_ENTRY_POINT', 'ODIN_NO_RTTI', 'ODIN_NO_TYPE_ASSERT', 'ODIN_OPTIMIZATION_MODE',
	'ODIN_OS', 'ODIN_PLATFORM_SUBTARGET', 'ODIN_ROOT', 'ODIN_TEST', 'ODIN_VENDOR', 'ODIN_VERSION',
	'ODIN_WINDOWS_SUBSYSTEM'
})

lex:set_word_list(lexer.VARIABLE_BUILTIN, {
	'context', 'runtime'
})

return lex