PgTidy/pkg/lexer/lexer.go

package lexer

import "strings"

// opChars are the characters PostgreSQL allows in operator names.
const opChars = "+-*/<>=~!@#%^&|`?"

// opSpecial is the subset whose presence lets an operator end in + or -.
const opSpecial = "~!@#%^&|`?"

// Lex tokenizes src into a lossless token stream. The concatenation of every
// returned token's Text equals src. The final token is always EOF (empty Text).
func Lex(src string) []Token {
	l := &lexer{src: src, line: 1, col: 1}
	return l.run()
}

type lexer struct {
	src  string
	pos  int
	line int
	col  int
	out  []Token

	// position of the token currently being scanned, captured each iteration
	tokOff, tokLine, tokCol int
}

func (l *lexer) run() []Token {
	n := len(l.src)
	for l.pos < n {
		l.tokOff, l.tokLine, l.tokCol = l.pos, l.line, l.col
		start := l.pos
		c := l.src[l.pos]
		switch {
		case isSpace(c):
			l.scanWhile(isSpace)
			l.emit(Whitespace, start)
		case c == '-' && l.peek(1) == '-':
			l.scanLineComment()
			l.emit(LineComment, start)
		case c == '/' && l.peek(1) == '*':
			l.scanBlockComment()
			l.emit(BlockComment, start)
		case c == '\'':
			l.scanString('\'', false)
			l.emit(String, start)
		case c == '"':
			l.scanString('"', false)
			l.emit(QuotedIdent, start)
		case c == '$':
			l.scanDollar(start)
		case c == '(':
			l.advance(1)
			l.emit(LParen, start)
		case c == ')':
			l.advance(1)
			l.emit(RParen, start)
		case c == '[':
			l.advance(1)
			l.emit(LBracket, start)
		case c == ']':
			l.advance(1)
			l.emit(RBracket, start)
		case c == ',':
			l.advance(1)
			l.emit(Comma, start)
		case c == ';':
			l.advance(1)
			l.emit(Semicolon, start)
		case c == ':':
			// :: (cast), := (assign), or bare : (slice).
			if l.peek(1) == ':' || l.peek(1) == '=' {
				l.advance(2)
			} else {
				l.advance(1)
			}
			l.emit(Operator, start)
		case c == '.':
			if isDigit(l.peek(1)) {
				l.scanNumber()
				l.emit(Number, start)
			} else {
				l.advance(1)
				l.emit(Dot, start)
			}
		case isDigit(c):
			l.scanNumber()
			l.emit(Number, start)
		case isIdentStart(c):
			l.scanWord(start)
		case strings.IndexByte(opChars, c) >= 0:
			l.scanOperator()
			l.emit(Operator, start)
		default:
			l.advance(1)
			l.emit(Unknown, start)
		}
	}
	l.out = append(l.out, Token{Kind: EOF, Off: l.pos, Line: l.line, Col: l.col})
	return l.out
}

// scanWord handles identifiers and the typed-string prefixes E' B' X' U&' / U&".
func (l *lexer) scanWord(start int) {
	c := l.src[l.pos]
	switch c {
	case 'E', 'e':
		if l.peek(1) == '\'' {
			l.advance(1)
			l.scanString('\'', true)
			l.emit(EscapeString, start)
			return
		}
	case 'B', 'b':
		if l.peek(1) == '\'' {
			l.advance(1)
			l.scanString('\'', false)
			l.emit(BitString, start)
			return
		}
	case 'X', 'x':
		if l.peek(1) == '\'' {
			l.advance(1)
			l.scanString('\'', false)
			l.emit(HexString, start)
			return
		}
	case 'U', 'u':
		if l.peek(1) == '&' && (l.peek(2) == '\'' || l.peek(2) == '"') {
			q := l.peek(2)
			l.advance(2)
			l.scanString(q, false)
			l.emit(UnicodeString, start)
			return
		}
	}
	l.scanWhile(isIdentCont)
	l.emit(Ident, start)
}

// scanDollar handles dollar-quoted strings ($tag$...$tag$), positional
// parameters ($1), and a lone $.
func (l *lexer) scanDollar(start int) {
	if tag, ok := l.dollarTag(); ok {
		// Opening delimiter is "$tag$"; find the matching close.
		open := "$" + tag + "$"
		l.advance(len(open))
		if idx := strings.Index(l.src[l.pos:], open); idx >= 0 {
			l.advance(idx + len(open))
		} else {
			l.advance(len(l.src) - l.pos) // unterminated: consume to EOF
		}
		l.emit(DollarString, start)
		return
	}
	if isDigit(l.peek(1)) {
		l.advance(1)
		l.scanWhile(isDigit)
		l.emit(Param, start)
		return
	}
	l.advance(1)
	l.emit(Unknown, start)
}

// dollarTag checks whether the current $ begins a dollar-quote opening
// delimiter and, if so, returns the (possibly empty) tag.
func (l *lexer) dollarTag() (string, bool) {
	// src[pos] == '$'
	j := l.pos + 1
	n := len(l.src)
	if j < n && l.src[j] == '$' {
		return "", true // "$$"
	}
	k := j
	if k < n && isIdentStart(l.src[k]) {
		k++
		// Tag chars follow identifier rules but exclude '$' (which closes the tag).
		for k < n && isIdentCont(l.src[k]) && l.src[k] != '$' {
			k++
		}
		if k < n && l.src[k] == '$' {
			return l.src[j:k], true
		}
	}
	return "", false
}

func (l *lexer) scanLineComment() {
	// Consume until newline (exclusive); the newline is whitespace.
	for l.pos < len(l.src) && l.src[l.pos] != '\n' {
		l.advance(1)
	}
}

func (l *lexer) scanBlockComment() {
	l.advance(2) // /*
	depth := 1
	for l.pos < len(l.src) && depth > 0 {
		if l.src[l.pos] == '/' && l.peek(1) == '*' {
			l.advance(2)
			depth++
		} else if l.src[l.pos] == '*' && l.peek(1) == '/' {
			l.advance(2)
			depth--
		} else {
			l.advance(1)
		}
	}
}

// scanString consumes a quoted run beginning at the current quote character.
// A doubled quote (”) is an embedded quote; with backslash=true a backslash
// escapes the next byte (escape-string syntax).
func (l *lexer) scanString(quote byte, backslash bool) {
	l.advance(1) // opening quote
	for l.pos < len(l.src) {
		c := l.src[l.pos]
		if backslash && c == '\\' && l.pos+1 < len(l.src) {
			l.advance(2)
			continue
		}
		if c == quote {
			if l.peek(1) == quote {
				l.advance(2) // doubled quote
				continue
			}
			l.advance(1) // closing quote
			return
		}
		l.advance(1)
	}
}

func (l *lexer) scanNumber() {
	// Non-decimal integer literals: 0x.., 0o.., 0b..
	if l.src[l.pos] == '0' {
		switch l.peek(1) {
		case 'x', 'X', 'o', 'O', 'b', 'B':
			l.advance(2)
			l.scanWhile(isHexOrSep)
			return
		}
	}
	l.scanWhile(isDigitOrSep)
	if l.pos < len(l.src) && l.src[l.pos] == '.' {
		l.advance(1)
		l.scanWhile(isDigitOrSep)
	}
	if c := l.cur(); c == 'e' || c == 'E' {
		if p := l.peek(1); isDigit(p) || ((p == '+' || p == '-') && isDigit(l.peek(2))) {
			l.advance(1)
			if l.cur() == '+' || l.cur() == '-' {
				l.advance(1)
			}
			l.scanWhile(isDigit)
		}
	}
}

// scanOperator consumes a run of operator characters, applying PostgreSQL's
// rule that a multi-character operator may not end in + or - unless it also
// contains one of ~ ! @ # % ^ & | ` ?.
func (l *lexer) scanOperator() {
	start := l.pos
	hasSpecial := false
	for l.pos < len(l.src) {
		c := l.src[l.pos]
		if strings.IndexByte(opChars, c) < 0 {
			break
		}
		if c == '-' && l.peek(1) == '-' {
			break // comment start
		}
		if c == '/' && l.peek(1) == '*' {
			break // comment start
		}
		if strings.IndexByte(opSpecial, c) >= 0 {
			hasSpecial = true
		}
		l.advance(1)
	}
	if !hasSpecial {
		for l.pos-1 > start {
			last := l.src[l.pos-1]
			if last != '+' && last != '-' {
				break
			}
			l.backup()
		}
	}
}

// --- low-level helpers ---

func (l *lexer) cur() byte {
	if l.pos < len(l.src) {
		return l.src[l.pos]
	}
	return 0
}

func (l *lexer) peek(k int) byte {
	if l.pos+k < len(l.src) {
		return l.src[l.pos+k]
	}
	return 0
}

func (l *lexer) scanWhile(pred func(byte) bool) {
	for l.pos < len(l.src) && pred(l.src[l.pos]) {
		l.advance(1)
	}
}

// advance moves forward k bytes, maintaining line/col.
func (l *lexer) advance(k int) {
	for i := 0; i < k && l.pos < len(l.src); i++ {
		if l.src[l.pos] == '\n' {
			l.line++
			l.col = 1
		} else {
			l.col++
		}
		l.pos++
	}
}

// backup moves back one byte. Only used within scanOperator, which never
// spans a newline, so column bookkeeping is safe.
func (l *lexer) backup() {
	l.pos--
	l.col--
}

// emit appends a token covering src[start:pos], using the start position
// captured at the top of the scan loop.
func (l *lexer) emit(kind Kind, start int) {
	l.out = append(l.out, Token{
		Kind: kind,
		Text: l.src[start:l.pos],
		Off:  l.tokOff,
		Line: l.tokLine,
		Col:  l.tokCol,
	})
}

func isSpace(c byte) bool {
	switch c {
	case ' ', '\t', '\n', '\r', '\v', '\f':
		return true
	}
	return false
}

func isDigit(c byte) bool      { return c >= '0' && c <= '9' }
func isDigitOrSep(c byte) bool { return isDigit(c) || c == '_' }

func isHexOrSep(c byte) bool {
	return isDigit(c) || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F') || c == '_'
}

func isIdentStart(c byte) bool {
	return c == '_' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c >= 0x80
}

func isIdentCont(c byte) bool {
	return isIdentStart(c) || isDigit(c) || c == '$'
}