Swift: extract comments

AlexDenisov · AlexDenisov · commit 5e69adb0a968 · 2022-08-05T11:50:48.000+02:00
diff --git a/swift/codegen/schema.yml b/swift/codegen/schema.yml
@@ -34,6 +34,10 @@ Location:
   end_column: int
   _pragma: qltest_skip
 
+Comment:
+  _extends: Locatable
+  text: string
+
 Type:
   name: string
   canonical_type: Type
diff --git a/swift/extractor/SwiftExtractor.cpp b/swift/extractor/SwiftExtractor.cpp
@@ -104,11 +104,26 @@ static void extractDeclarations(const SwiftExtractorConfiguration& config,
   trap.emit(unknownFileEntry);
   trap.emit(unknownLocationEntry);
 
+  std::vector<swift::Token> comments;
+  if (primaryFile && primaryFile->getBufferID().hasValue()) {
+    auto& sourceManager = compiler.getSourceMgr();
+    auto tokens = swift::tokenize(compiler.getInvocation().getLangOptions(), sourceManager,
+                                  primaryFile->getBufferID().getValue());
+    for (auto& token : tokens) {
+      if (token.getKind() == swift::tok::comment) {
+        comments.push_back(token);
+      }
+    }
+  }
+
   SwiftVisitor visitor(compiler.getSourceMgr(), trap, module, primaryFile);
   auto topLevelDecls = getTopLevelDecls(module, primaryFile);
   for (auto decl : topLevelDecls) {
     visitor.extract(decl);
   }
+  for (auto& comment : comments) {
+    visitor.extract(comment);
+  }
 }
 
 static std::unordered_set<std::string> collectInputFilenames(swift::CompilerInstance& compiler) {
diff --git a/swift/extractor/infra/SwiftDispatcher.h b/swift/extractor/infra/SwiftDispatcher.h
@@ -3,6 +3,7 @@
 #include <swift/AST/SourceFile.h>
 #include <swift/Basic/SourceManager.h>
 #include <llvm/Support/FileSystem.h>
+#include <swift/Parse/Token.h>
 
 #include "swift/extractor/trap/TrapLabelStore.h"
 #include "swift/extractor/trap/TrapDomain.h"
@@ -242,6 +243,12 @@ class SwiftDispatcher {
     return false;
   }
 
+  void emitComment(swift::Token& comment) {
+    CommentsTrap entry{trap.createLabel<CommentTag>(), comment.getRawText().str()};
+    trap.emit(entry);
+    attachLocation(comment.getRange().getStart(), comment.getRange().getEnd(), entry.id);
+  }
+
  private:
   void attachLocation(swift::SourceLoc start,
                       swift::SourceLoc end,
diff --git a/swift/extractor/visitors/SwiftVisitor.h b/swift/extractor/visitors/SwiftVisitor.h
@@ -17,6 +17,7 @@ class SwiftVisitor : private SwiftDispatcher {
   void extract(const T& entity) {
     fetchLabel(entity);
   }
+  void extract(swift::Token& comment) { emitComment(comment); }
 
  private:
   void visit(swift::Decl* decl) override { declVisitor.visit(decl); }
diff --git a/swift/ql/lib/codeql/swift/elements.qll b/swift/ql/lib/codeql/swift/elements.qll
@@ -1,6 +1,7 @@
 // generated by codegen/codegen.py
 import codeql.swift.elements.AstNode
 import codeql.swift.elements.Callable
+import codeql.swift.elements.Comment
 import codeql.swift.elements.Element
 import codeql.swift.elements.File
 import codeql.swift.elements.Locatable
diff --git a/swift/ql/lib/codeql/swift/elements/Comment.qll b/swift/ql/lib/codeql/swift/elements/Comment.qll
@@ -0,0 +1,6 @@
+private import codeql.swift.generated.Comment
+
+class Comment extends CommentBase {
+  /** toString */
+  override string toString() { result = getText() }
+}
diff --git a/swift/ql/lib/codeql/swift/generated/Comment.qll b/swift/ql/lib/codeql/swift/generated/Comment.qll
@@ -0,0 +1,8 @@
+// generated by codegen/codegen.py
+import codeql.swift.elements.Locatable
+
+class CommentBase extends @comment, Locatable {
+  override string getAPrimaryQlClass() { result = "Comment" }
+
+  string getText() { comments(this, result) }
+}
diff --git a/swift/ql/lib/swift.dbscheme b/swift/ql/lib/swift.dbscheme
@@ -35,6 +35,7 @@ files(
 @locatable =
   @argument
 | @ast_node
+| @comment
 | @condition_element
 | @if_config_clause
 ;
@@ -54,6 +55,11 @@ locations(
   int end_column: int ref
 );
 
+comments(
+  unique int id: @comment,
+  string text: string ref
+);
+
 @type =
   @any_function_type
 | @any_generic_type
diff --git a/swift/ql/test/extractor-tests/comments/comments.expected b/swift/ql/test/extractor-tests/comments/comments.expected
@@ -0,0 +1,4 @@
+| comments.swift:1:1:2:1 | // Single line comment\n |
+| comments.swift:3:1:6:3 | /*\n  Multiline\n  comment\n*/ |
+| comments.swift:8:1:9:1 | /// Single line doc comment\n |
+| comments.swift:10:1:13:3 | /**\n  Multiline\n  doc comment\n*/ |
diff --git a/swift/ql/test/extractor-tests/comments/comments.ql b/swift/ql/test/extractor-tests/comments/comments.ql
@@ -0,0 +1,4 @@
+import swift
+
+from Comment c
+select c
diff --git a/swift/ql/test/extractor-tests/comments/comments.swift b/swift/ql/test/extractor-tests/comments/comments.swift
@@ -0,0 +1,13 @@
+// Single line comment
+
+/*
+  Multiline
+  comment
+*/
+
+/// Single line doc comment
+
+/**
+  Multiline
+  doc comment
+*/
diff --git a/swift/ql/test/extractor-tests/generated/Comment/MISSING_SOURCE.txt b/swift/ql/test/extractor-tests/generated/Comment/MISSING_SOURCE.txt
@@ -0,0 +1,4 @@
+// generated by codegen/codegen.py
+
+After a swift source file is added in this directory and codegen/codegen.py is run again, test queries
+will appear and this file will be deleted

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,7 @@ class SwiftVisitor : private SwiftDispatcher {`
`17`	`17`	`void extract(const T& entity) {`
`18`	`18`	`fetchLabel(entity);`
`19`	`19`	`}`
	`20`	`+ void extract(swift::Token& comment) { emitComment(comment); }`
`20`	`21`
`21`	`22`	`private:`
`22`	`23`	`void visit(swift::Decl* decl) override { declVisitor.visit(decl); }`