compiler: Add Ability to Annotate Preprocessed Source

Signed-off-by: Will Hawkins <hawkinsw@obs.cr>

Sources/Common/Formatter.swift

@@ -75,6 +75,20 @@ public struct FormatterPlain: Formattable {
 
 }
 
+public struct FormatterDelimited: Formattable {
+  let start: String
+  let end: String
+
+  public init(_ start: String, _ end: String) {
+    self.start = start
+    self.end = end
+  }
+
+  public func formatWithStyle(_ value: String, _ style: Style) -> String {
+    return self.start + value + self.end
+  }
+}
+
 public struct FormatterAnsi: Formattable {
 
   public init() {}

Sources/Common/SourceCode.swift

@@ -143,14 +143,57 @@ public struct SourceCode {
     self.locations = locations
   }
 
-  public func getSource() -> String {
-    return self.code
-  }
-
   public func getManager() -> SourceManager {
     return self.manager
   }
 
+  static func do_annotate(
+    _ contents: String, _ manager: SourceManager, _ locations: FileSourceLocation,
+    _ formatter: Formattable, _ style: Style
+  ) -> String {
+    var result = ""
+
+    // Keep track of the start of any gap between nested locations.
+    var gap_start = contents.startIndex
+
+    // contents are devoid of any preceding source code, but the locations do not know that. So,
+    // when we use a range from locations we must adjust appropriately.
+    let offset = locations.location.range.lowerBound
+
+    for nested in locations.getNestedLocations() {
+      let nested_start = contents.index(
+        contents.startIndex, offsetBy: nested.location.range.lowerBound - offset)
+      let nested_end = contents.index(
+        contents.startIndex, offsetBy: nested.location.range.upperBound - offset)
+
+      // Add in any gap.
+      result += contents[gap_start..<nested_start]
+
+      // Handle this range (recursively)
+      result += do_annotate(
+        "\(contents[nested_start ..< nested_end])", manager, nested, formatter, style)
+
+      // Adjust where the next gap will start.
+      gap_start = nested_end
+    }
+
+    // Is there anything left?
+    let remainder = contents[gap_start...]
+
+    return formatter.formatWithStyle(result + remainder, style)
+  }
+
+  public func getSource(
+    annotated: Bool = false, formatter: Formattable = FormatterDelimited("<", ">"),
+    initialStyle: Style = Style(StyleColor.Red)
+  ) -> String {
+    if annotated {
+      return SourceCode.do_annotate(
+        self.code, self.manager, self.locations, formatter, initialStyle)
+    }
+    return self.code
+  }
+
   public func getLocations() -> FileSourceLocation {
     return self.locations
   }

TestData/Sources/annotate-parser-state.p4

@@ -0,0 +1,9 @@
+    state start {
+        Testing ts;
+        ts.yesno = true;
+        ts.count = 5;
+        transition select (ts.count == 5) {
+            true: accept;
+            false: reject;
+        };
+    }

TestData/Sources/annotate-parser.p4

@@ -0,0 +1,3 @@
+parser main_parser() {
+#include <annotate-parser-state.p4>
+}

TestData/Sources/annotate-struct-body.p4

@@ -0,0 +1,2 @@
+    bool yesno;
+    int count;

TestData/Sources/annotate.p4

@@ -0,0 +1,4 @@
+struct Testing {
+#include <annotate-struct-body.p4>
+};
+#include <annotate-parser.p4>

Tests/p4rseTests/SupportTests/SourceCode.swift

@@ -156,6 +156,35 @@ import TreeSitterP4
       == SourceLocation(48..<166))
 }
 
+@Test func test_preprocessor_nested_includes_annotated_source() async throws {
+  let sm = SourceManager(["./TestData/Sources/"])
+  let prep = SourceCodePreprocessor(sm)
+  let file = FilePath.init(stringLiteral: "./TestData/Sources/annotate.p4")
+
+  let expected = """
+    <struct Testing {
+    <    bool yesno;
+        int count;>
+    };
+    <parser main_parser() {
+    <    state start {
+            Testing ts;
+            ts.yesno = true;
+            ts.count = 5;
+            transition select (ts.count == 5) {
+                true: accept;
+                false: reject;
+            };
+        }
+    >
+    }>>
+    """
+
+  #expect(#RequireOkResult(prep.preprocess(file)))
+  let source = try! (#UseOkResult(prep.preprocess(file)))
+  #expect(source.getSource(annotated: true) == expected)
+}
+
 
 @Test func test_source_location_contains() async throws {
   let outer = SourceLocation(0..<500)