Improve pipeline expressions lexer

1 files changed, 27 insertions, 15 deletions

diff --git a/lib/gitlab/ci/pipeline/expression/lexer.rb b/lib/gitlab/ci/pipeline/expression/lexer.rb
index 067ad5fd312..cabb91d0dd7 100644
--- a/lib/gitlab/ci/pipeline/expression/lexer.rb
+++ b/lib/gitlab/ci/pipeline/expression/lexer.rb
@@ -3,6 +3,8 @@ module Gitlab
     module Pipeline
       module Expression
         class Lexer
+          include ::Gitlab::Utils::StrongMemoize
+
           LEXEMES = [
             Expression::Lexeme::Variable,
             Expression::Lexeme::String,
@@ -10,34 +12,44 @@ module Gitlab
             Expression::Lexeme::Equals
           ].freeze
 
-          MAX_CYCLES = 5
           SyntaxError = Class.new(Statement::StatementError)
 
+          MAX_TOKENS = 100
+
           def initialize(statement)
             @scanner = StringScanner.new(statement)
-            @tokens = []
           end
 
-          def tokens
-            return @tokens if @tokens.any?
+          def tokens(max: MAX_TOKENS)
+            strong_memoize(:tokens) { tokenize(max) }
+          end
+
+          def lexemes
+            tokens.map(&:to_lexeme)
+          end
+
+          private
 
-            MAX_CYCLES.times do
-              LEXEMES.each do |lexeme|
-                @scanner.skip(/\s+/) # ignore whitespace
+          def tokenize(max_tokens)
+            tokens = []
 
-                lexeme.scan(@scanner).tap do |token|
-                  @tokens.push(token) if token.present?
+            max_tokens.times do
+              @scanner.skip(/\s+/) # ignore whitespace
+
+              return tokens if @scanner.eos?
+
+              lexeme = LEXEMES.find do |type|
+                type.scan(@scanner).tap do |token|
+                  tokens.push(token) if token.present?
                 end
+              end
 
-                return @tokens if @scanner.eos?
+              unless lexeme.present?
+                raise Lexer::SyntaxError, 'Unknown lexeme found!'
               end
             end
 
-            raise Lexer::SyntaxError unless @scanner.eos?
-          end
-
-          def lexemes
-            tokens.map(&:to_lexeme)
+            raise Lexer::SyntaxError, 'Too many tokens!'
           end
         end
       end