From ad40fd023a25ad75e1b1a3a04244ee72079e6f89 Mon Sep 17 00:00:00 2001
From: Boik <qazbnm456@gmail.com>
Date: Tue, 10 Jul 2018 11:10:36 +0800
Subject: [PATCH] lexer: higher coverage of the 'SET Syntax'

---
 parser/lexer.go      | 8 +++-----
 parser/lexer_test.go | 6 +++++-
 parser/misc.go       | 4 ++++
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/parser/lexer.go b/parser/lexer.go
index a2522d0cd3123..8ce6c01379752 100644
--- a/parser/lexer.go
+++ b/parser/lexer.go
@@ -435,8 +435,8 @@ func startWithAt(s *Scanner) (tok int, pos Pos, lit string) {
 		} else {
 			tok = int('@')
 		}
-	} else if isIdentChar(ch1) {
-		s.r.incAsLongAs(isIdentChar)
+	} else if isUserVarChar(ch1) {
+		s.r.incAsLongAs(isUserVarChar)
 		tok, lit = singleAtIdentifier, s.r.data(&pos)
 	} else if ch1 == '@' {
 		s.r.inc()
@@ -452,10 +452,8 @@ func startWithAt(s *Scanner) (tok int, pos Pos, lit string) {
 		}
 		s.r.incAsLongAs(isIdentChar)
 		tok, lit = doubleAtIdentifier, s.r.data(&pos)
-	} else if ch1 == ' ' || ch1 == ':' || ch1 == '=' {
-		tok, lit = singleAtIdentifier, s.r.data(&pos)
 	} else {
-		tok = int('@')
+		tok, lit = singleAtIdentifier, s.r.data(&pos)
 	}
 	return
 }
diff --git a/parser/lexer_test.go b/parser/lexer_test.go
index 4cde88755e09f..a720d4acf814f 100644
--- a/parser/lexer_test.go
+++ b/parser/lexer_test.go
@@ -71,10 +71,14 @@ func (s *testLexerSuite) TestAtLeadingIdentifier(c *C) {
 		{"@", singleAtIdentifier},
 		{"@''", singleAtIdentifier},
 		{"@1", singleAtIdentifier},
+		{"@.1_", singleAtIdentifier},
+		{"@-1.", singleAtIdentifier},
+		{"@~", singleAtIdentifier},
 		{"@$", singleAtIdentifier},
 		{"@a_3cbbc", singleAtIdentifier},
 		{"@`a_3cbbc`", singleAtIdentifier},
-		{"@-3cbbc", int('@')},
+		{"@-3cbbc", singleAtIdentifier},
+		{"@!3cbbc", singleAtIdentifier},
 		{"@@global.test", doubleAtIdentifier},
 		{"@@session.test", doubleAtIdentifier},
 		{"@@local.test", doubleAtIdentifier},
diff --git a/parser/misc.go b/parser/misc.go
index c953b01466007..e0183805bccc0 100644
--- a/parser/misc.go
+++ b/parser/misc.go
@@ -36,6 +36,10 @@ func isIdentExtend(ch rune) bool {
 	return ch >= 0x80 && ch <= '\uffff'
 }
 
+func isUserVarChar(ch rune) bool {
+	return isLetter(ch) || isDigit(ch) || ch == '_' || ch == '$' || ch == '.' || isIdentExtend(ch)
+}
+
 type trieNode struct {
 	childs [256]*trieNode
 	token  int