Merge pull request tensor-compiler#334 from roastduck/fix-parse

Fix -s argument parsing in the command line tool
rohany · Dec 15, 2020 · dcbfdc3 · dcbfdc3
2 parents dd62216 + 1b27e3c
commit dcbfdc3
Showing 1 changed file with 8 additions and 12 deletions.
diff --git a/tools/taco.cpp b/tools/taco.cpp
@@ -814,23 +814,19 @@ int main(int argc, char* argv[]) {
     }
     else if ("-s" == argName) {  
       setSchedule = true;       
-      bool insideCall = false; 
-      bool parsingExpr = false; 
+      int parenthesesCnt = 0;
 
-      std::replace_if(argValue.begin(), argValue.end(), [&insideCall, &parsingExpr](char c) {
+      std::replace_if(argValue.begin(), argValue.end(), [&parenthesesCnt](char c) {
         if (c == '(') {
-          if (insideCall) {
-            parsingExpr = true; // need to handle precompute case specially
-          } else {
-            insideCall = true; 
-            return true; 
+          if (parenthesesCnt++ == 0) { // '(' for a call
+            return true;
           }
         } else if (c == ',') {
-          return !parsingExpr; 
+          return parenthesesCnt <= 1;
         } else if (c == ')') {
-          bool previous = parsingExpr; 
-          parsingExpr = false; 
-          return !previous; 
+          if (--parenthesesCnt == 0) { // ')' for a call
+            return true;
+          }
         }
         return false; 
       }, ' ');