IDE: implement an IDE action to collect types of all expressions in a source file.

This is libIDE side implementation for collecting all type information in a source
file. When several expression share the same source range, we always report the
type of the outermost expression.

rdar://35199889

Author: nkcsgexi

include/swift/Sema/IDETypeChecking.h

@@ -161,6 +161,28 @@ namespace swift {
     bool shouldPrintRequirement(ExtensionDecl *ED, StringRef Req);
     bool hasMergeGroup(MergeGroupKind Kind);
   };
+
+  /// Reported type for an expression. This expression is represented by offset
+  /// length in the source buffer;
+  struct ExpressionTypeInfo {
+
+    /// The start of the expression;
+    uint32_t offset;
+
+    /// The length of the expression;
+    uint32_t length;
+
+    /// The start of the printed type in a separately given string buffer.
+    uint32_t typeOffset;
+
+    /// The length of the printed type
+    uint32_t typeLength;
+  };
+
+  /// Collect type information for every expression in \c SF; all types will
+  /// be printed to \c OS.
+  ArrayRef<ExpressionTypeInfo> collectExpressionType(SourceFile &SF,
+    std::vector<ExpressionTypeInfo> &scratch, llvm::raw_ostream &OS);
 }
 
 #endif

lib/IDE/IDETypeChecking.cpp

@@ -23,6 +23,8 @@
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/Sema/IDETypeChecking.h"
+#include "swift/IDE/SourceEntityWalker.h"
+#include "swift/Parse/Lexer.h"
 
 using namespace swift;
 
@@ -577,3 +579,82 @@ collectDefaultImplementationForProtocolMembers(ProtocolDecl *PD,
   for (auto *IP : PD->getInheritedProtocols())
     HandleMembers(IP->getMembers());
 }
+
+/// This walker will traverse the AST and report types for every expression.
+class ExpressionTypeCollector: public SourceEntityWalker {
+  SourceManager &SM;
+  unsigned int BufferId;
+  std::vector<ExpressionTypeInfo> &Results;
+
+  // This is to where we print all types.
+  llvm::raw_ostream &OS;
+
+  // Map from a printed type to the offset in OS where the type starts.
+  llvm::StringMap<uint32_t> TypeOffsets;
+
+  // This keeps track of whether we have a type reported for a given
+  // [offset, length].
+  llvm::DenseMap<unsigned, llvm::DenseSet<unsigned>> AllPrintedTypes;
+
+  bool shouldReport(unsigned Offset, unsigned Length, Expr *E) {
+    // We shouldn't report null types.
+    if (E->getType().isNull())
+      return false;
+
+    // If we have already reported types for this source range, we shouldn't
+    // report again. This makes sure we always report the outtermost type of
+    // several overlapping expressions.
+    auto &Bucket = AllPrintedTypes[Offset];
+    return Bucket.find(Length) == Bucket.end();
+  }
+
+  // Find an existing offset in the type buffer otherwise print the type to
+  // the buffer.
+  uint32_t getTypeOffsets(StringRef PrintedType) {
+    auto It = TypeOffsets.find(PrintedType);
+    if (It == TypeOffsets.end()) {
+      TypeOffsets[PrintedType] = OS.tell();
+      OS << PrintedType;
+    }
+    return TypeOffsets[PrintedType];
+  }
+
+public:
+  ExpressionTypeCollector(SourceFile &SF, std::vector<ExpressionTypeInfo> &Results,
+    llvm::raw_ostream &OS): SM(SF.getASTContext().SourceMgr),
+                            BufferId(*SF.getBufferID()),
+                            Results(Results), OS(OS) {}
+  bool walkToExprPre(Expr *E) override {
+    if (E->getSourceRange().isInvalid())
+      return true;
+    CharSourceRange Range =
+      Lexer::getCharSourceRangeFromSourceRange(SM, E->getSourceRange());
+    unsigned Offset = SM.getLocOffsetInBuffer(Range.getStart(), BufferId);
+    unsigned Length = Range.getByteLength();
+    if (!shouldReport(Offset, Length, E))
+      return true;
+    // Print the type to a temporary buffer.
+    SmallString<64> Buffer;
+    {
+      llvm::raw_svector_ostream OS(Buffer);
+      E->getType()->getRValueType()->reconstituteSugar(true)->print(OS);
+    }
+
+    // Add the type information to the result list.
+    Results.push_back({Offset, Length, getTypeOffsets(Buffer.str()),
+      static_cast<uint32_t>(Buffer.size())});
+
+    // Keep track of that we have a type reported for this range.
+    AllPrintedTypes[Offset].insert(Length);
+    return true;
+  }
+};
+
+ArrayRef<ExpressionTypeInfo>
+swift::collectExpressionType(SourceFile &SF,
+                             std::vector<ExpressionTypeInfo> &Scratch,
+                             llvm::raw_ostream &OS) {
+  ExpressionTypeCollector Walker(SF, Scratch, OS);
+  Walker.walk(SF);
+  return Scratch;
+}

test/IDE/Inputs/ExprType.swift

@@ -0,0 +1,22 @@
+
+func foo() -> Int  { return 1 }
+
+func bar(f: Float) -> Float { return bar(f: 1) }
+
+protocol P {}
+
+func fooP(_ p: P) { fooP(p) }
+
+class C {}
+
+func ArrayC(_ a: [C]) {
+	_ = a.count
+	_ = a.description.count.advanced(by: 1).description
+}
+
+struct S {
+  let val = 4
+}
+func DictS(_ a: [Int: S]) {
+  _ = a[2]?.val.advanced(by: 1).byteSwapped
+}

test/IDE/expr_type.swift

@@ -0,0 +1,8 @@
+// RUN: %target-swift-ide-test -print-expr-type -source-filename %S/Inputs/ExprType.swift -swift-version 5 | %FileCheck %s
+
+// CHECK: func foo() -> Int { return <expr type:"Int">1</expr> }
+// CHECK: func bar(f: Float) -> Float { return <expr type:"Float"><expr type:"(Float) -> Float">bar</expr><expr type:"(f: Float)">(f: <expr type:"Float">1</expr>)</expr></expr> }
+// CHECK: func fooP(_ p: P) { <expr type:"()"><expr type:"(P) -> ()">fooP</expr><expr type:"(P)">(<expr type:"P">p</expr>)</expr></expr> }
+// CHECK: <expr type:"()"><expr type:"Int">_</expr> = <expr type:"Int"><expr type:"[C]">a</expr>.count</expr></expr>
+// CHECK: <expr type:"()"><expr type:"String">_</expr> = <expr type:"String"><expr type:"Int"><expr type:"(Int) -> Int"><expr type:"Int"><expr type:"String"><expr type:"[C]">a</expr>.description</expr>.count</expr>.<expr type:"(Int) -> (Int) -> Int">advanced</expr></expr><expr type:"(by: Int)">(by: <expr type:"Int">1</expr>)</expr></expr>.description</expr></expr>
+// CHECK: <expr type:"()"><expr type:"Int?">_</expr> = <expr type:"Int?"><expr type:"Int"><expr type:"(Int) -> Int"><expr type:"Int"><expr type:"S"><expr type:"S?"><expr type:"[Int : S]">a</expr><expr type:"(Int)">[<expr type:"Int">2</expr>]</expr></expr>?</expr>.val</expr>.<expr type:"(Int) -> (Int) -> Int">advanced</expr></expr><expr type:"(by: Int)">(by: <expr type:"Int">1</expr>)</expr></expr>.byteSwapped</expr></expr>

tools/swift-ide-test/swift-ide-test.cpp

@@ -95,6 +95,7 @@ enum class ActionType {
   PrintLocalTypes,
   PrintTypeInterface,
   PrintIndexedSymbols,
+  PrintExpressionTypes,
   TestCreateCompilerInvocation,
   CompilerInvocationFromModule,
   GenerateModuleAPIDescription,
@@ -228,6 +229,9 @@ Action(llvm::cl::desc("Mode:"), llvm::cl::init(ActionType::None),
            clEnumValN(ActionType::TypeContextInfo,
 	                    "type-context-info",
                       "Perform expression context info analysis"),
+           clEnumValN(ActionType::PrintExpressionTypes,
+                      "print-expr-type",
+                      "Print types for all expressions in the file"),
            clEnumValN(ActionType::ConformingMethodList,
 	                    "conforming-methods",
                       "Perform conforming method analysis for expression")));
@@ -1709,6 +1713,57 @@ static int doPrintAST(const CompilerInvocation &InitInvok,
   return EXIT_SUCCESS;
 }
 
+static int doPrintExpressionTypes(const CompilerInvocation &InitInvok,
+                                  StringRef SourceFilename) {
+  CompilerInvocation Invocation(InitInvok);
+  Invocation.getFrontendOptions().InputsAndOutputs.addPrimaryInputFile(SourceFilename);
+  CompilerInstance CI;
+
+  // Display diagnostics to stderr.
+  PrintingDiagnosticConsumer PrintDiags;
+  CI.addDiagnosticConsumer(&PrintDiags);
+  if (CI.setup(Invocation))
+    return EXIT_FAILURE;
+  CI.performSema();
+  std::vector<ExpressionTypeInfo> Scratch;
+
+  // Buffer for where types will be printed.
+  llvm::SmallString<256> TypeBuffer;
+  llvm::raw_svector_ostream OS(TypeBuffer);
+  SourceFile &SF = *CI.getPrimarySourceFile();
+  auto Source = SF.getASTContext().SourceMgr.getRangeForBuffer(*SF.getBufferID()).str();
+  std::vector<std::pair<unsigned, std::string>> SortedTags;
+
+  // Collect all tags of expressions.
+  for (auto R: collectExpressionType(*CI.getPrimarySourceFile(), Scratch, OS)) {
+    SortedTags.push_back({R.offset,
+      (llvm::Twine("<expr type:\"") + TypeBuffer.str().substr(R.typeOffset,
+                                                  R.typeLength) + "\">").str()});
+    SortedTags.push_back({R.offset + R.length, "</expr>"});
+  }
+  // Sort tags by offset.
+  std::stable_sort(SortedTags.begin(), SortedTags.end(),
+    [](std::pair<unsigned, std::string> T1, std::pair<unsigned, std::string> T2) {
+      return T1.first < T2.first;
+  });
+
+  ArrayRef<std::pair<unsigned, std::string>> SortedTagsRef = SortedTags;
+  unsigned Cur = 0;
+  do {
+    // Print tags that are due at current offset.
+    while(!SortedTagsRef.empty() && SortedTagsRef.front().first == Cur) {
+      llvm::outs() << SortedTagsRef.front().second;
+      SortedTagsRef = SortedTagsRef.drop_front();
+    }
+    auto Start = Cur;
+    // Change current offset to the start offset of next tag.
+    Cur = SortedTagsRef.empty() ? Source.size() : SortedTagsRef.front().first;
+    // Print the source before next tag.
+    llvm::outs() << Source.substr(Start, Cur - Start);
+  } while(!SortedTagsRef.empty());
+  return EXIT_SUCCESS;
+}
+
 static int doPrintLocalTypes(const CompilerInvocation &InitInvok,
                              const std::vector<std::string> ModulesToPrint) {
   using NodeKind = Demangle::Node::Kind;
@@ -3353,6 +3408,12 @@ int main(int argc, char *argv[]) {
                                   options::CodeCompletionDiagnostics);
     break;
 
+  case ActionType::PrintExpressionTypes:
+    ExitCode = doPrintExpressionTypes(InitInvok,
+                                      options::SourceFilename);
+    break;
+
+
   case ActionType::ConformingMethodList:
     if (options::CodeCompletionToken.empty()) {
       llvm::errs() << "token name required\n";