Fix exp-time parse backtrack by saving if '[' starts Array or Dictionary

include/swift/Parse/Parser.h

@@ -348,6 +348,75 @@ class Parser {
   /// Current syntax parsing context where call backs should be directed to.
   SyntaxParsingContext *SyntaxContext;
 
+  // Save the results of discovering whether a given '[' is the start of
+  // an Array or a Dictionary. Remembering this prevents
+  // `parseExprCollection()` from having to parse the first sub-expr in
+  // square brackets {2^n} times, where n is # of nested square brackets.
+  class ArrayOrDictStartLocCache {
+    typedef unsigned SquareBracketLoc;
+
+    llvm::DenseMap<SquareBracketLoc, bool> IsDictMap;
+    int UseCount;
+    unsigned BufferID;
+
+  public:
+    ArrayOrDictStartLocCache() : UseCount(0), BufferID(INT_MAX) {}
+
+    struct CachedVal {
+      bool IsCached;
+      bool IsDict;
+      CachedVal(bool IsCached, bool IsDict)
+          : IsCached(IsCached), IsDict(IsDict) {}
+    };
+
+    // Handle used by all the Parser::parseExprCollection() calls on the stack
+    class RAII {
+      Parser &P;
+      /// Translate a ParserPosition into something we can use as a map key.
+      SquareBracketLoc TranslatePos(SourceLoc Loc) {
+        // SourceLoc Loc = Pos.PreviousLoc;
+        assert(
+            P.BracketCache.BufferID ==
+                P.SourceMgr.findBufferContainingLoc(Loc) &&
+            "square-bracket sub-expressions are assumed not to cross files!");
+        return P.SourceMgr.getLocOffsetInBuffer(Loc, P.BracketCache.BufferID);
+      }
+
+    public:
+      RAII(Parser &P) : P(P) {
+        auto &cache = P.BracketCache;
+        if (cache.UseCount == 0) {
+          cache.BufferID = P.SourceMgr.findBufferContainingLoc(
+              P.getParserPosition().PreviousLoc);
+        }
+        ++cache.UseCount;
+      }
+      ~RAII() {
+        auto &cache = P.BracketCache;
+        --cache.UseCount;
+        if (cache.UseCount == 0) { // The topmost speculative parse has ended
+          cache.IsDictMap.shrink_and_clear();
+          cache.BufferID = INT_MAX;
+        }
+      }
+
+      void setIsDictStarting(SourceLoc Loc, bool IsDict) {
+        SquareBracketLoc SBLoc = TranslatePos(Loc);
+        P.BracketCache.IsDictMap[SBLoc] = IsDict;
+      }
+
+      CachedVal whatStartsAt(SourceLoc Loc) {
+        SquareBracketLoc SBLoc = TranslatePos(Loc);
+        auto i = P.BracketCache.IsDictMap.find(SBLoc);
+        if (i == P.BracketCache.IsDictMap.end())
+          return CachedVal(false, false);
+        else
+          return CachedVal(true, i->second);
+      }
+    };
+  };
+  ArrayOrDictStartLocCache BracketCache;
+
 public:
   Parser(unsigned BufferID, SourceFile &SF, DiagnosticEngine* LexerDiags,
          SILParserTUStateBase *SIL,

lib/Parse/ParseExpr.cpp

@@ -3382,6 +3382,13 @@ ParserResult<Expr> Parser::parseExprCollection() {
   }
 
   bool ParseDict;
+  ArrayOrDictStartLocCache::RAII SquareCacheRAII(*this);
+
+  ArrayOrDictStartLocCache::CachedVal v =
+      SquareCacheRAII.whatStartsAt(LSquareLoc);
+  if (v.IsCached) {
+    ParseDict = v.IsDict; // true for dict, false for array
+  } else                  // First time parsing this or neither dict nor array
   {
     BacktrackingScope Scope(*this);
     auto HasDelayedDecl = State->hasDelayedDecl();
@@ -3400,6 +3407,7 @@ ParserResult<Expr> Parser::parseExprCollection() {
       State->takeDelayedDeclState();
     // If we have a ':', this is a dictionary literal.
     ParseDict = Tok.is(tok::colon);
+    SquareCacheRAII.setIsDictStarting(LSquareLoc, ParseDict);
   }
 
   if (ParseDict) {

test/Parse/structure_no_overflow.swift

@@ -35,3 +35,16 @@ let c = ({({({({ ({({({([ ({({({({ ({({({([  ({({({({ ({({({([ ({({({({ ({({({({
         })})})}) })})})}) ])})})}) })})})})  ])})})}) })})})}) ])})})}) })})})}) 
         })})})}) })})})}) ])})})}) })})})})  ])})})}) })})})}) ])})})}) })})})}) 
         })})})}) })})})}) ])})})}) })})})})  ])})})}) })})})}) ])})})}) })})})}) 
+
+// Test that we can have the allowed # of square brackets.
+// Note in bug SR-9220 this gave no error, but failed to parse in reasonable time.
+// CHECK-NOT: error
+let d = [[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[ [[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[
+        [[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[ [[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[
+        [[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[ [[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[
+        [[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[ [[[[[[[[[[[[[[[[[[[[[[[[[[[[[[[
+         1
+         ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]] ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]
+        ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]] ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]
+        ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]] ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]
+        ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]] ]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]]