
import clay.common.*;
import clay.diagnostics.*;
import maybe.*;
import parsing.combinators.wrapper.*;



//
// data definitions
//

enum TokenKind (
    SENTINEL,
    PUNCTUATION,
    KEYWORD,
    IDENTIFIER,
    STRING_LITERAL,
    CHAR_LITERAL,
    INT_LITERAL,
    FLOAT_LITERAL,
    SPACE,
    COMMENT,
    LLVM,
);

alias DataRange = Range[SizeT];

record Token(kind:TokenKind, subKind:UInt8, range:DataRange);

[I | Integer?(I)]
overload Token(kind:TokenKind, subKind:I) =
    Token(kind, UInt8(subKind), DataRange(SizeT(0), SizeT(0)));

overload Token(kind:TokenKind) = Token(kind, 0);

record SourceTokens = referenceType(
    file: SourceFile,
    vector: Vector[Token],
);



//
// punctuations and keywords
//

private var punctuations = array(
    "<--", "...", "=>", "->",
    "==", "!=", "<=", ">=",
    "<", ">",
    "+=", "-=", "*=", "/=", "%=",
    "+", "-", "*", "/", "%", "=", "&", "^", "|",
    "(", ")", "[", "]", "{", "}",
    ":", ";", ",", ".", "#"
);

private var keywords = array(
    "public", "private", "import", "as",
    "record", "variant", "instance",
    "procedure", "overload", "external", "alias",
    "static", "callbyname", "lvalue", "rvalue",
    "inline", "enum", "var", "ref", "forward",
    "and", "or", "not", "new",
    "if", "else", "goto", "return", "while",
    "switch", "case", "default", "break", "continue", "for", "in",
    "true", "false", "try", "catch", "throw",
    "__FILE__", "__LINE__", "__COLUMN__",
);

punctuationIndex(s) : Int {
    for (i, x in enumerated(punctuations)) {
        if (x == s)
            return Int(i);
    }
    return -1;
}

keywordIndex(s) : Int {
    for (i, x in enumerated(keywords)) {
        if (x == s)
            return Int(i);
    }
    return -1;
}



//
// lexer input
//

record LexerInput(
    file : SourceFile,
    current : SizeT,
    maxCurrent : SizeT,
);

overload LexerInput(file : SourceFile) =
    LexerInput(file, SizeT(0), SizeT(0));

overload iterator(x:LexerInput) = x;

overload hasNext?(x:LexerInput) = (x.current < size(x.file.data));

overload next(x:LexerInput) {
    var c = x.file.data[x.current];
    x.current += 1;
    return c;
}

overload assign(ref dest:LexerInput, ref src:LexerInput) {
    if (dest.file == src.file) {
        var destMax = max(dest.maxCurrent, dest.current);
        var srcMax = max(src.maxCurrent, src.current);
        dest.maxCurrent = max(destMax, srcMax);
        dest.current = src.current;
    }
    else {
        dest.file = src.file;
        dest.current = src.current;
        dest.maxCurrent = src.maxCurrent;
    }
}



//
// combinator specialization
//

private var combinators = Combinators[LexerInput]();

private var Parser? = combinators.Parser?;
private var Parsers? = combinators.Parsers?;
private var ParseResult = combinators.ParseResult;

private var nextChar = combinators.nextToken;

private var condition = combinators.condition;
private var literal = combinators.literal;
private var literalString = combinators.literalString;
private var optional = combinators.optional;
private var choice = combinators.choice;
private var sequence = combinators.sequence;
private var sequenceAs = combinators.sequenceAs;
private var modify = combinators.modify;
private var repeated = combinators.repeated;
private var repeatedAs = combinators.repeatedAs;
private var zeroOrMore = combinators.zeroOrMore;
private var oneOrMore = combinators.oneOrMore;
private var zeroOrMoreAs = combinators.zeroOrMoreAs;
private var oneOrMoreAs = combinators.oneOrMoreAs;



//
// utility
//

private alias Dummy = Char;
private alias dummy = 'T';

private seq(forward ..parsers) = sequenceAs(dummy, ..parsers);

private zeroPlus(forward ..parsers) = zeroPlus(seq(..parsers));
overload zeroPlus(forward parser) = zeroOrMoreAs(dummy, parser);

private onePlus(forward ..parsers) = onePlus(seq(..parsers));
overload onePlus(forward parser) = oneOrMoreAs(dummy, parser);



//
// bounds combinator
//

[Parser]
private bounds(parser:Parser) {
    return input => {
        alias T = ParseResult(Parser);
        var start = input.current;
        return maybe(parser(input),
            x -> {
                x.range = DataRange(start, input.current);
                return Maybe(x);
            },
            () => nothing(T),
        );
    };
}



//
// char predicates
//

private decimalDigit?(c) = ((c >= '0') and (c <= '9'));
private hexDigit?(c) = decimalDigit?(c) or
                       ((c >= 'a') and (c <= 'f')) or
                       ((c >= 'A') and (c <= 'F'));
private octalDigit?(c) = ((c >= '0') and (c <= '7'));

private alpha?(c) = ((c >= 'a') and (c <= 'z')) or
                    ((c >= 'A') and (c <= 'Z'));



//
// identifiers and keywords
//

private identChar1?(c) = alpha?(c) or (c == '_');
private identChar2?(c) = identChar1?(c) or decimalDigit?(c) or (c == '?');

private var identChar1 = condition(identChar1?);
private var identChar2 = condition(identChar2?);
private var identStr = seq(identChar1, zeroPlus(identChar2));

private lexIdentifierOrKeyword(input) : Maybe[Token] {
    var start = input.current;
    if (not nothing?(identStr(input))) {
        var data = sliced(input.file.data, start, input.current);
        var i = keywordIndex(data);
        if (i >= 0)
            return Maybe(Token(KEYWORD, i));
        return Maybe(Token(IDENTIFIER));
    }
    return nothing(Token);
}



//
// punctuation
//

private lexPunctuation(input) {
    for (i, sym in enumerated(punctuations)) {
        var parser = literalString(sym);
        var result = parser(input);
        if (not nothing?(result))
            return Maybe(Token(PUNCTUATION, i));
    }
    return nothing(Token);
}



//
// hexDigit, octalDigit, decimalDigit
//

private var hexDigit = condition(hexDigit?);
private var octalDigit = condition(octalDigit?);
private var decimalDigit = condition(decimalDigit?);



//
// char literal, string literal
//

private var hexEscape = seq('x', hexDigit, hexDigit);
private var octalEscape = seq(octalDigit, octalDigit, octalDigit);

private var escapeSeq = choice(
    'n', 'r', 't', 'a', 'b', 'f', 'v', '\\', '\'', '\"',
     hexEscape, octalEscape,
);

private var escapeChar2 = seq('\\', escapeSeq);
private escapeChar(input) = escapeChar2(input);

private var charChar = choice(
    condition(c => (c != '\\' and c != '\'')),
    escapeChar,
);

private var strChar = choice(
    condition(c => (c != '\\' and c != '\"')),
    escapeChar,
);

private var lexCharLiteral = modify(
    x => Token(CHAR_LITERAL),
    seq('\'', charChar, '\''),
);

private var lexStringLiteral = modify(
    x => Token(STRING_LITERAL),
    seq('\"', zeroPlus(strChar), '\"'),
);



//
// int literal
//

private var hexInt = seq(
    "0x", hexDigit, zeroPlus(choice('_', hexDigit))
);

private var decimalInt = seq(
    decimalDigit, zeroPlus(choice('_', decimalDigit))
);

private var sign = choice('+', '-');

private var lexIntLiteral = modify(
    x => Token(INT_LITERAL),
    seq(optional(sign), choice(hexInt, decimalInt)),
);



//
// float literal
//

private var exponent = seq(choice('e','E'), optional(sign), decimalInt);
private var fractional = seq('.', decimalInt);

private var floatTail = choice(
    seq(fractional, optional(exponent)),
    exponent,
);

private var lexFloatLiteral = modify(
    x => Token(FLOAT_LITERAL),
    seq(optional(sign), decimalInt, floatTail),
);



//
// space
//

private space?(c) {
    for (x in " \t\n\r\f\x0B") {
        if (c == x)
            return true;
    }
    return false;
}

private var lexSpace = modify(
    x => Token(SPACE),
    onePlus(condition(space?)),
);



//
// comments
//

private var notNewline = condition(c => (c != '\n'));

private var lexLineComment = modify(
    x => Token(COMMENT),
    seq("//", zeroPlus(notNewline), optional('\n')),
);

private blockCommentTail(input) {
    var saved = input;
    var lastWasStar = false;
    while (hasNext?(input)) {
        var c = next(input);
        if (lastWasStar and (c == '/'))
            return Maybe(dummy);
        lastWasStar = (c == '*');
    }
    input = saved;
    return nothing(Dummy);
}

private var lexBlockComment = modify(
    x => Token(COMMENT),
    seq("/*", blockCommentTail)
);



//
// inline llvm
//

private llvmBodyItem(input) : Maybe[Dummy] = llvmBodyItem2(input);

private var llvmComment = seq(';', zeroPlus(notNewline), '\n');

private var llvmStringChar = choice(
    seq('\\', nextChar),
    condition(c => (c != '\"')),
);
private var llvmString = seq('\"', zeroPlus(llvmStringChar), '\"');


private var llvmBraces = seq('{', zeroPlus(llvmBodyItem), '}');

private var llvmBodyItem2 = choice(
    llvmComment,
    llvmString,
    llvmBraces,
    condition(c => (c != '}')),
);

private var lexLLVM = modify(
    x => Token(LLVM),
    seq("__llvm__", zeroPlus(condition(space?)), llvmBraces),
);



//
// lex one token
//

private var lexToken = bounds(choice(
    lexSpace,
    lexLineComment,
    lexBlockComment,
    lexPunctuation,
    lexLLVM,
    lexIdentifierOrKeyword,
    lexCharLiteral,
    lexStringLiteral,
    lexFloatLiteral,
    lexIntLiteral,
));



//
// LexerError
//

record LexerError(
    file: SourceFile,
    where: SizeT
);

instance ClayError (LexerError);

overload displayError(e:LexerError) {
    errorWithLocation(e.file, e.where, "invalid token");
}



//
// tokenize
//

private var lexSentinel = bounds(input => Maybe(Token(SENTINEL)));

private var tokenize1 = modify(
    x => {
        var v = move(x.0);
        push(v, move(x.1));
        return move(v);
    },
    sequence(zeroOrMore(lexToken), lexSentinel),
);

tokenizeAll(file) {
    var input = LexerInput(file);
    var tokensData = require(tokenize1(input));
    if (input.current != size(file.data))
        throw LexerError(file, input.maxCurrent);
    return SourceTokens(file, move(tokensData));
}

tokenize(file) {
    var tokens = tokenizeAll(file);
    tokens.vector = Vector(filtered(
        x => (x.kind != SPACE) and (x.kind != COMMENT),
        tokens.vector,
    ));
    return move(tokens);
}
