Feat/exact custom token

lexer/lexer.go

@@ -20,7 +20,7 @@ type Lexer struct {
 	peeks  []token.Token
 	isEOF  bool
 
-	customs map[string]struct{}
+	customs map[string]token.TokenType
 }
 
 func New(r io.Reader, opts ...OptionFunc) *Lexer {
@@ -29,10 +29,9 @@ func New(r io.Reader, opts ...OptionFunc) *Lexer {
 		r:       bufio.NewReader(r),
 		line:    1,
 		buffer:  new(bytes.Buffer),
-		customs: map[string]struct{}{},
+		file:    o.Filename,
+		customs: o.Customs,
 	}
-	l.file = o.Filename
-	l.customs = o.Customs
 	l.readChar()
 	return l
 }
@@ -41,9 +40,9 @@ func NewFromString(input string, opts ...OptionFunc) *Lexer {
 	return New(strings.NewReader(input), opts...)
 }
 
-func (l *Lexer) RegisterCustomTokens(tokens ...string) {
-	for i := range tokens {
-		l.customs[tokens[i]] = struct{}{}
+func (l *Lexer) RegisterCustomTokens(tokenMap map[string]token.TokenType) {
+	for k, v := range tokenMap {
+		l.customs[k] = v
 	}
 }
 
@@ -360,9 +359,9 @@ func (l *Lexer) NextToken() token.Token {
 				}
 			default:
 				t.Literal = literal
-				// If custom token found, mark as CUSTOM
-				if _, ok := l.customs[literal]; ok {
-					t.Type = token.CUSTOM
+				// If custom token found, use it
+				if custom, ok := l.customs[literal]; ok {
+					t.Type = custom
 				} else {
 					t.Type = token.LookupIdent(t.Literal)
 				}

lexer/lexer_test.go

@@ -681,10 +681,12 @@ func TestPeekToken(t *testing.T) {
 
 func TestCustomToken(t *testing.T) {
 	input := `describe foo {}`
-	l := NewFromString(input, WithCustomTokens("describe"))
+	l := NewFromString(input, WithCustomTokens(map[string]token.TokenType{
+		"describe": token.Custom("DESCRIBE"),
+	}))
 
 	expects := []token.Token{
-		{Type: token.CUSTOM, Literal: "describe", Line: 1, Position: 1},
+		{Type: token.TokenType("DESCRIBE"), Literal: "describe", Line: 1, Position: 1},
 		{Type: token.IDENT, Literal: "foo", Line: 1, Position: 10},
 		{Type: token.LEFT_BRACE, Literal: "{", Line: 1, Position: 14},
 		{Type: token.RIGHT_BRACE, Literal: "}", Line: 1, Position: 15},

lexer/option.go

@@ -1,10 +1,12 @@
 package lexer
 
+import "github.com/ysugimoto/falco/token"
+
 type OptionFunc func(o *Option)
 
 type Option struct {
 	Filename string
-	Customs  map[string]struct{}
+	Customs  map[string]token.TokenType
 	// more field if exists
 }
 
@@ -14,18 +16,18 @@ func WithFile(filename string) OptionFunc {
 	}
 }
 
-func WithCustomTokens(idents ...string) OptionFunc {
+func WithCustomTokens(tokenMap map[string]token.TokenType) OptionFunc {
 	return func(o *Option) {
-		for i := range idents {
-			o.Customs[idents[i]] = struct{}{}
+		for k, v := range tokenMap {
+			o.Customs[k] = v
 		}
 	}
 }
 
 func collect(opts []OptionFunc) *Option {
 	o := &Option{
 		Filename: "",
-		Customs:  make(map[string]struct{}),
+		Customs:  make(map[string]token.TokenType),
 	}
 
 	for i := range opts {

parser/custom_parser.go

@@ -2,15 +2,17 @@ package parser
 
 import (
 	"github.com/ysugimoto/falco/ast"
+	"github.com/ysugimoto/falco/token"
 )
 
 type CustomParser interface {
-	Literal() string
+	Ident() string
+	Token() token.TokenType
 	Parse(*Parser) (ast.CustomStatement, error)
 }
 
 func (p *Parser) ParseCustomToken() (ast.CustomStatement, error) {
-	v, ok := p.customParsers[p.curToken.Token.Literal]
+	v, ok := p.customParsers[p.curToken.Token.Type]
 	if !ok {
 		return nil, UnexpectedToken(p.curToken)
 	}

parser/custom_parser_test.go

@@ -49,9 +49,13 @@ func (d *DescribeStatement) String() string {
 
 type DescribeParser struct{}
 
-func (d *DescribeParser) Literal() string {
+func (d *DescribeParser) Ident() string {
 	return "describe"
 }
+func (d *DescribeParser) Token() token.TokenType {
+	return token.Custom("DESCRIBE")
+}
+
 func (d *DescribeParser) Parse(p *Parser) (ast.CustomStatement, error) {
 	stmt := &DescribeStatement{
 		Meta: p.curToken,
@@ -75,7 +79,7 @@ func (d *DescribeParser) Parse(p *Parser) (ast.CustomStatement, error) {
 				return nil, errors.WithStack(err)
 			}
 			stmt.Subroutines = append(stmt.Subroutines, sub)
-		case token.CUSTOM:
+		case token.TokenType("BEFORE_EACH"):
 			cs, err := p.ParseCustomToken()
 			if err != nil {
 				return nil, errors.WithStack(err)
@@ -125,9 +129,12 @@ func (b *BeforeEachStatement) String() string {
 
 type BeforeEachParser struct{}
 
-func (b *BeforeEachParser) Literal() string {
+func (b *BeforeEachParser) Ident() string {
 	return "before_each"
 }
+func (b *BeforeEachParser) Token() token.TokenType {
+	return token.Custom("BEFORE_EACH")
+}
 func (b *BeforeEachParser) Parse(p *Parser) (ast.CustomStatement, error) {
 	stmt := &BeforeEachStatement{
 		Meta: p.curToken,

parser/option.go

@@ -5,7 +5,7 @@ type ParserOption func(p *Parser)
 func WithCustomParser(cps ...CustomParser) ParserOption {
 	return func(p *Parser) {
 		for i := range cps {
-			p.customParsers[cps[i].Literal()] = cps[i]
+			p.customParsers[cps[i].Token()] = cps[i]
 		}
 	}
 }

parser/parser.go

@@ -57,13 +57,13 @@ type Parser struct {
 	prefixParsers  map[token.TokenType]prefixParser
 	infixParsers   map[token.TokenType]infixParser
 	postfixParsers map[token.TokenType]postfixParser
-	customParsers  map[string]CustomParser
+	customParsers  map[token.TokenType]CustomParser
 }
 
 func New(l *lexer.Lexer, opts ...ParserOption) *Parser {
 	p := &Parser{
 		l:             l,
-		customParsers: make(map[string]CustomParser),
+		customParsers: make(map[token.TokenType]CustomParser),
 	}
 	for i := range opts {
 		opts[i](p)
@@ -72,11 +72,11 @@ func New(l *lexer.Lexer, opts ...ParserOption) *Parser {
 	p.registerExpressionParsers()
 
 	// Register custom lexer tokens for each custom parsers
-	var literals []string
-	for literal := range p.customParsers {
-		literals = append(literals, literal)
+	customTokenMap := make(map[string]token.TokenType)
+	for _, v := range p.customParsers {
+		customTokenMap[v.Ident()] = v.Token()
 	}
-	l.RegisterCustomTokens(literals...)
+	l.RegisterCustomTokens(customTokenMap)
 
 	p.NextToken()
 	p.NextToken()
@@ -241,10 +241,12 @@ func (p *Parser) Parse() (ast.Statement, error) {
 		stmt, err = p.ParsePenaltyboxDeclaration()
 	case token.RATECOUNTER:
 		stmt, err = p.ParseRatecounterDeclaration()
-	case token.CUSTOM:
-		stmt, err = p.ParseCustomToken()
 	default:
-		err = UnexpectedToken(p.curToken)
+		if custom, ok := p.customParsers[p.curToken.Token.Type]; ok {
+			stmt, err = custom.Parse(p)
+		} else {
+			err = UnexpectedToken(p.curToken)
+		}
 	}
 
 	if err != nil {

parser/statement_parser.go

@@ -61,8 +61,6 @@ func (p *Parser) ParseStatement() (ast.Statement, error) {
 		stmt, err = p.ParseBreakStatement()
 	case token.FALLTHROUGH:
 		stmt, err = p.ParseFallthroughStatement()
-	case token.CUSTOM:
-		stmt, err = p.ParseCustomToken()
 	case token.IDENT:
 		// Check if the current ident is a function call
 		if p.PeekTokenIs(token.LEFT_PAREN) {
@@ -76,7 +74,11 @@ func (p *Parser) ParseStatement() (ast.Statement, error) {
 			}
 		}
 	default:
-		err = UnexpectedToken(p.curToken)
+		if custom, ok := p.customParsers[p.curToken.Token.Type]; ok {
+			stmt, err = custom.Parse(p)
+		} else {
+			err = UnexpectedToken(p.curToken)
+		}
 	}
 	if err != nil {
 		return nil, errors.WithStack(err)

tester/syntax/describe.go

@@ -57,9 +57,12 @@ func (d *DescribeStatement) String() string {
 // Custome parser implementation for "describe" keyword
 type DescribeParser struct{}
 
-func (d *DescribeParser) Literal() string {
+func (d *DescribeParser) Ident() string {
 	return "describe"
 }
+func (d *DescribeParser) Token() token.TokenType {
+	return token.Custom("DESCRIBE")
+}
 func (d *DescribeParser) Parse(p *parser.Parser) (ast.CustomStatement, error) {
 	stmt := &DescribeStatement{
 		Meta:    p.CurToken(),
@@ -86,38 +89,39 @@ func (d *DescribeParser) Parse(p *parser.Parser) (ast.CustomStatement, error) {
 			}
 			stmt.Subroutines = append(stmt.Subroutines, sub)
 			p.NextToken()
-		case token.CUSTOM:
-			cs, err := p.ParseCustomToken()
-			if err != nil {
-				return nil, errors.WithStack(err)
-			}
-			if t, ok := cs.(*HookStatement); ok {
-				switch {
-				case strings.HasPrefix(t.keyword, "before"):
-					if _, ok := stmt.Befores[t.keyword]; ok {
-						return nil, &parser.ParseError{
-							Token:   t.GetMeta().Token,
-							Message: fmt.Sprintf("%s hook is duplicated", cs.Literal()),
+		default:
+			if hookParser, ok := hookParsers[tok.Token.Type]; ok {
+				hook, err := hookParser.Parse(p)
+				if err != nil {
+					return nil, errors.WithStack(err)
+				}
+				if t, ok := hook.(*HookStatement); ok {
+					switch {
+					case strings.HasPrefix(t.keyword, "before"):
+						if _, ok := stmt.Befores[t.keyword]; ok {
+							return nil, &parser.ParseError{
+								Token:   t.GetMeta().Token,
+								Message: fmt.Sprintf("%s hook is duplicated", hook.Literal()),
+							}
 						}
-					}
-					stmt.Befores[t.keyword] = t
-				case strings.HasPrefix(t.keyword, "after"):
-					if _, ok := stmt.Afters[t.keyword]; ok {
-						return nil, &parser.ParseError{
-							Token:   t.GetMeta().Token,
-							Message: fmt.Sprintf("%s hook is duplicated", cs.Literal()),
+						stmt.Befores[t.keyword] = t
+					case strings.HasPrefix(t.keyword, "after"):
+						if _, ok := stmt.Afters[t.keyword]; ok {
+							return nil, &parser.ParseError{
+								Token:   t.GetMeta().Token,
+								Message: fmt.Sprintf("%s hook is duplicated", hook.Literal()),
+							}
 						}
+						stmt.Afters[t.keyword] = t
 					}
-					stmt.Afters[t.keyword] = t
+					p.NextToken()
+					continue
+				}
+				return nil, &parser.ParseError{
+					Token:   p.CurToken().Token,
+					Message: fmt.Sprintf("%s statement could not be placed inside describe", hook.Literal()),
 				}
-				p.NextToken()
-				continue
-			}
-			return nil, &parser.ParseError{
-				Token:   p.CurToken().Token,
-				Message: fmt.Sprintf("%s statement could not be placed inside describe", cs.Literal()),
 			}
-		default:
 			return nil, parser.UnexpectedToken(p.CurToken())
 		}
 	}

tester/syntax/hooks.go

@@ -2,6 +2,7 @@ package syntax
 
 import (
 	"bytes"
+	"strings"
 
 	"github.com/pkg/errors"
 	"github.com/ysugimoto/falco/ast"
@@ -42,9 +43,12 @@ type HookParser struct {
 	keyword string
 }
 
-func (h *HookParser) Literal() string {
+func (h *HookParser) Ident() string {
 	return h.keyword
 }
+func (h *HookParser) Token() token.TokenType {
+	return token.Custom(strings.ToUpper(h.keyword))
+}
 func (h *HookParser) Parse(p *parser.Parser) (ast.CustomStatement, error) {
 	stmt := &HookStatement{
 		Meta:    p.CurToken(),
@@ -61,3 +65,24 @@ func (h *HookParser) Parse(p *parser.Parser) (ast.CustomStatement, error) {
 
 	return stmt, nil
 }
+
+var hookParsers = map[token.TokenType]*HookParser{
+	token.Custom("BEFORE_RECV"):    {keyword: "before_recv"},
+	token.Custom("BEFORE_HASH"):    {keyword: "before_hash"},
+	token.Custom("BEFORE_HIT"):     {keyword: "before_hit"},
+	token.Custom("BEFORE_MISS"):    {keyword: "before_miss"},
+	token.Custom("BEFORE_PASS"):    {keyword: "before_pass"},
+	token.Custom("BEFORE_FETCH"):   {keyword: "before_fetch"},
+	token.Custom("BEFORE_ERROR"):   {keyword: "before_error"},
+	token.Custom("BEFORE_DELIVER"): {keyword: "before_deliver"},
+	token.Custom("BEFORE_LOG"):     {keyword: "before_log"},
+	token.Custom("AFTER_RECV"):     {keyword: "after_recv"},
+	token.Custom("AFTER_HASH"):     {keyword: "after_hash"},
+	token.Custom("AFTER_HIT"):      {keyword: "after_hit"},
+	token.Custom("AFTER_MISS"):     {keyword: "after_miss"},
+	token.Custom("AFTER_PASS"):     {keyword: "after_pass"},
+	token.Custom("AFTER_FETCH"):    {keyword: "after_fetch"},
+	token.Custom("AFTER_ERROR"):    {keyword: "after_error"},
+	token.Custom("AFTER_DELIVER"):  {keyword: "after_deliver"},
+	token.Custom("AFTER_LOG"):      {keyword: "after_log"},
+}