From 01e06397c138a42516f65e16782eb069651b29fa Mon Sep 17 00:00:00 2001
From: Connie Zhu <connieyzhu@google.com>
Date: Sun, 11 Aug 2024 20:41:21 +0000
Subject: [PATCH] [llvm-lit] Initial draft of curly brace implementation

This is a draft of the parsing implementation to support curly brace
syntax in lit's internal shell.
---
 llvm/utils/lit/lit/ShCommands.py |  5 +++--
 llvm/utils/lit/lit/ShUtil.py     | 36 +++++++++++++++++++++++++++++---
 llvm/utils/lit/lit/TestRunner.py |  2 +-
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/llvm/utils/lit/lit/ShCommands.py b/llvm/utils/lit/lit/ShCommands.py
index 68655a41d7934..392e05aaabdcd 100644
--- a/llvm/utils/lit/lit/ShCommands.py
+++ b/llvm/utils/lit/lit/ShCommands.py
@@ -92,14 +92,15 @@ def toShell(self, file, pipefail=False):
 
 
 class Seq:
-    def __init__(self, lhs, op, rhs):
+    def __init__(self, lhs, op, rhs, seq_type):
         assert op in (";", "&", "||", "&&")
         self.op = op
         self.lhs = lhs
         self.rhs = rhs
+        self.type = seq_type
 
     def __repr__(self):
-        return "Seq(%r, %r, %r)" % (self.lhs, self.op, self.rhs)
+        return "Seq(%r, %r, %r, %r)" % (self.lhs, self.op, self.rhs, self.type)
 
     def __eq__(self, other):
         if not isinstance(other, Seq):
diff --git a/llvm/utils/lit/lit/ShUtil.py b/llvm/utils/lit/lit/ShUtil.py
index fa13167cad1be..b835abe5a48ec 100644
--- a/llvm/utils/lit/lit/ShUtil.py
+++ b/llvm/utils/lit/lit/ShUtil.py
@@ -157,7 +157,11 @@ def lex_one_token(self):
         lex_one_token - Lex a single 'sh' token."""
 
         c = self.eat()
+        if c == "{" or c == "}":
+            return (c,)
         if c == ";":
+            if self.maybe_eat("}") or (self.maybe_eat(" ") and self.maybe_eat("}")):
+                return("}",)
             return (c,)
         if c == "|":
             if self.maybe_eat("|"):
@@ -200,6 +204,8 @@ def __init__(self, data, win32Escapes=False, pipefail=False):
         self.data = data
         self.pipefail = pipefail
         self.tokens = ShLexer(data, win32Escapes=win32Escapes).lex()
+        self.brace_stack = []
+        self.brace_dict = {'{': '}'}
 
     def lex(self):
         for item in self.tokens:
@@ -255,18 +261,42 @@ def parse_pipeline(self):
             self.lex()
             commands.append(self.parse_command())
         return Pipeline(commands, negate, self.pipefail)
+            
+    
+    # {echo foo; echo bar;} && echo hello
+    # echo hello && {echo foo; echo bar}
+    
+    def parse(self, seq_type):
+        lhs = None
+        if isinstance(self.look(), tuple):
+            brace = self.lex()
+            self.brace_stack.append(brace)
+            if brace[0] == '{':
+                lhs = self.parse(('{', '}'))
+            else:
+                raise ValueError("syntax error near unexpected token %r" % brace[0])
 
-    def parse(self):
-        lhs = self.parse_pipeline()
+        else:
+            lhs = self.parse_pipeline()
 
         while self.look():
             operator = self.lex()
             assert isinstance(operator, tuple) and len(operator) == 1
 
+            if operator == self.brace_dict[self.brace_stack.peek()]:
+                break
+
             if not self.look():
                 raise ValueError("missing argument to operator %r" % operator[0])
 
             # FIXME: Operator precedence!!
-            lhs = Seq(lhs, operator[0], self.parse_pipeline())
+            if isinstance(self.look(), tuple):
+                lhs = self.parse(('{', '}'))
+            else: 
+                lhs = Seq(lhs, operator[0], self.parse_pipeline(), seq_type)
+            seq_type = None
+        
+        if not stack.empty():
+            raise ValueError("missing token to %r" % stack.peek())
 
         return lhs
diff --git a/llvm/utils/lit/lit/TestRunner.py b/llvm/utils/lit/lit/TestRunner.py
index da7fa86fd3917..e5b484dbc62e8 100644
--- a/llvm/utils/lit/lit/TestRunner.py
+++ b/llvm/utils/lit/lit/TestRunner.py
@@ -1056,7 +1056,7 @@ def executeScriptInternal(
             ln = command
         try:
             cmds.append(
-                ShUtil.ShParser(ln, litConfig.isWindows, test.config.pipefail).parse()
+                ShUtil.ShParser(ln, litConfig.isWindows, test.config.pipefail).parse(None)
             )
         except:
             raise ScriptFatal(