[SourceKit/CodeFormat] Re-work and improve the indentation implementation

include/swift/AST/Expr.h

@@ -566,11 +566,14 @@ class alignas(8) Expr {
 /// typically this will have an ErrorType.
 class ErrorExpr : public Expr {
   SourceRange Range;
+  Expr *OriginalExpr;
 public:
-  ErrorExpr(SourceRange Range, Type Ty = Type())
-    : Expr(ExprKind::Error, /*Implicit=*/true, Ty), Range(Range) {}
+  ErrorExpr(SourceRange Range, Type Ty = Type(), Expr *OriginalExpr = nullptr)
+    : Expr(ExprKind::Error, /*Implicit=*/true, Ty), Range(Range),
+      OriginalExpr(OriginalExpr) {}
 
   SourceRange getSourceRange() const { return Range; }
+  Expr *getOriginalExpr() const { return OriginalExpr; }
 
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::Error;

include/swift/AST/Stmt.h

@@ -321,6 +321,7 @@ class alignas(8) PoundAvailableInfo final :
   friend TrailingObjects;
 
   SourceLoc PoundLoc;
+  SourceLoc LParenLoc;
   SourceLoc RParenLoc;
 
   // The number of queries tail allocated after this object.
@@ -338,17 +339,18 @@ class alignas(8) PoundAvailableInfo final :
   /// This is filled in by Sema.
   VersionRange VariantAvailableRange;
 
-  PoundAvailableInfo(SourceLoc PoundLoc, ArrayRef<AvailabilitySpec *> queries,
-                     SourceLoc RParenLoc)
-   : PoundLoc(PoundLoc), RParenLoc(RParenLoc), NumQueries(queries.size()),
-     AvailableRange(VersionRange::empty()),
+  PoundAvailableInfo(SourceLoc PoundLoc, SourceLoc LParenLoc,
+                     ArrayRef<AvailabilitySpec *> queries, SourceLoc RParenLoc)
+   : PoundLoc(PoundLoc), LParenLoc(LParenLoc), RParenLoc(RParenLoc),
+     NumQueries(queries.size()), AvailableRange(VersionRange::empty()),
      VariantAvailableRange(VersionRange::empty()) {
     std::uninitialized_copy(queries.begin(), queries.end(),
                             getTrailingObjects<AvailabilitySpec *>());
   }
 
 public:
   static PoundAvailableInfo *create(ASTContext &ctx, SourceLoc PoundLoc,
+                                    SourceLoc LParenLoc,
                                     ArrayRef<AvailabilitySpec *> queries,
                                     SourceLoc RParenLoc);
 
@@ -357,6 +359,9 @@ class alignas(8) PoundAvailableInfo final :
                               NumQueries);
   }
 
+  SourceLoc getLParenLoc() const { return LParenLoc; }
+  SourceLoc getRParenLoc() const { return RParenLoc; }
+
   SourceLoc getStartLoc() const { return PoundLoc; }
   SourceLoc getEndLoc() const;
   SourceLoc getLoc() const { return PoundLoc; }
@@ -771,6 +776,7 @@ class WhileStmt : public LabeledConditionalStmt {
 
   SourceLoc getStartLoc() const { return getLabelLocOrKeywordLoc(WhileLoc); }
   SourceLoc getEndLoc() const { return Body->getEndLoc(); }
+  SourceLoc getWhileLoc() const { return WhileLoc; }
 
   Stmt *getBody() const { return Body; }
   void setBody(Stmt *s) { Body = s; }
@@ -792,9 +798,10 @@ class RepeatWhileStmt : public LabeledStmt {
                   getDefaultImplicitFlag(implicit, RepeatLoc),
                   LabelInfo),
       RepeatLoc(RepeatLoc), WhileLoc(WhileLoc), Body(Body), Cond(Cond) {}
-  
+
   SourceLoc getStartLoc() const { return getLabelLocOrKeywordLoc(RepeatLoc); }
   SourceLoc getEndLoc() const;
+  SourceLoc getRepeatLoc() const { return RepeatLoc; }
 
   Stmt *getBody() const { return Body; }
   void setBody(Stmt *s) { Body = s; }
@@ -819,6 +826,7 @@ class ForEachStmt : public LabeledStmt {
   Pattern *Pat;
   SourceLoc InLoc;
   Expr *Sequence;
+  SourceLoc WhereLoc;
   Expr *WhereExpr = nullptr;
   BraceStmt *Body;
 
@@ -831,12 +839,12 @@ class ForEachStmt : public LabeledStmt {
 
 public:
   ForEachStmt(LabeledStmtInfo LabelInfo, SourceLoc ForLoc, Pattern *Pat,
-              SourceLoc InLoc, Expr *Sequence, Expr *WhereExpr, BraceStmt *Body,
-              Optional<bool> implicit = None)
+              SourceLoc InLoc, Expr *Sequence, SourceLoc WhereLoc,
+              Expr *WhereExpr, BraceStmt *Body, Optional<bool> implicit = None)
     : LabeledStmt(StmtKind::ForEach, getDefaultImplicitFlag(implicit, ForLoc),
                   LabelInfo),
       ForLoc(ForLoc), Pat(nullptr), InLoc(InLoc), Sequence(Sequence),
-      WhereExpr(WhereExpr), Body(Body) {
+      WhereLoc(WhereLoc), WhereExpr(WhereExpr), Body(Body) {
     setPattern(Pat);
   }
 
@@ -864,6 +872,9 @@ class ForEachStmt : public LabeledStmt {
 
   /// getInLoc - Retrieve the location of the 'in' keyword.
   SourceLoc getInLoc() const { return InLoc; }
+
+  /// getWhereLoc - Retrieve the location of the 'where' keyword.
+  SourceLoc getWhereLoc() const { return WhereLoc; }
 
   /// getPattern - Retrieve the pattern describing the iteration variables.
   /// These variables will only be visible within the body of the loop.

lib/AST/ASTWalker.cpp

@@ -327,6 +327,31 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     return false;
   }
 
+  bool visitParamDecl(ParamDecl *P) {
+    // Don't walk into the type if the decl is implicit, or if the type is
+    // implicit.
+    if (!P->isImplicit()) {
+      if (auto *repr = P->getTypeRepr()) {
+        if (doIt(repr)) {
+          return true;
+        }
+      }
+    }
+    if (auto *E = P->getStructuralDefaultExpr()) {
+      auto res = doIt(E);
+      if (!res) return true;
+      P->setDefaultExpr(res, /*isTypeChecked*/ (bool)res->getType());
+    }
+
+    if (!Walker.shouldWalkAccessorsTheOldWay()) {
+      for (auto *AD : P->getAllAccessors())
+        if (doIt(AD))
+          return true;
+    }
+
+    return false;
+  }
+
   bool visitSubscriptDecl(SubscriptDecl *SD) {
     bool WalkGenerics = visitGenericParamListIfNeeded(SD);
 
@@ -1142,22 +1167,6 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       // Walk each parameter's decl and typeloc and default value.
       if (doIt(P))
         return true;
-
-      // Don't walk into the type if the decl is implicit, or if the type is
-      // implicit.
-      if (!P->isImplicit()) {
-        if (auto *repr = P->getTypeRepr()) {
-          if (doIt(repr)) {
-            return true;
-          }
-        }
-      }
-
-      if (auto *E = P->getStructuralDefaultExpr()) {
-        auto res = doIt(E);
-        if (!res) return true;
-        P->setDefaultExpr(res, /*isTypeChecked*/ (bool)res->getType());
-      }
     }
 
     return Walker.walkToParameterListPost(PL);

lib/AST/Stmt.cpp

@@ -305,17 +305,21 @@ bool CatchStmt::isSyntacticallyExhaustive() const {
 
 PoundAvailableInfo *PoundAvailableInfo::create(ASTContext &ctx,
                                                SourceLoc PoundLoc,
+                                               SourceLoc LParenLoc,
                                        ArrayRef<AvailabilitySpec *> queries,
                                                      SourceLoc RParenLoc) {
   unsigned size = totalSizeToAlloc<AvailabilitySpec *>(queries.size());
   void *Buffer = ctx.Allocate(size, alignof(PoundAvailableInfo));
-  return ::new (Buffer) PoundAvailableInfo(PoundLoc, queries, RParenLoc);
+  return ::new (Buffer) PoundAvailableInfo(PoundLoc, LParenLoc, queries,
+                                           RParenLoc);
 }
 
 SourceLoc PoundAvailableInfo::getEndLoc() const {
   if (RParenLoc.isInvalid()) {
     if (NumQueries == 0) {
-      return PoundLoc;
+      if (LParenLoc.isInvalid())
+        return PoundLoc;
+      return LParenLoc;
     }
     return getQueries()[NumQueries - 1]->getSourceRange().End;
   }