Fix AllowUnsafeNonSealedKeyType annotation placement for same-line function params

Summary:
When a function parameter needing thrift.AllowUnsafeNonSealedKeyType
shared a line with the function signature, the annotation was incorrectly
placed before the entire function line. Add annotate_function_param() that
detects this case and breaks the line after the preceding ( or , separator,
placing the annotation on its own indented line before the parameter.

Reviewed By: hchokshi

Differential Revision: D95133891

fbshipit-source-id: 7bbca99608804a5346569854dedde374a29e99ca

Author: aristidisp

third-party/thrift/src/thrift/compiler/codemod/annotate_allow_unsafe_non_sealed_key_type.cc

@@ -87,7 +87,7 @@ class AnnotateAllowUnsafeNonSealedKeyType final {
       for (const t_field& param : func.params().fields()) {
         const t_type& type = param.type().deref();
         if (needs_annotation(type, param)) {
-          annotate_node(param.type().src_range().begin);
+          annotate_function_param(param, func);
           any_annotated = true;
         }
       }
@@ -138,6 +138,58 @@ class AnnotateAllowUnsafeNonSealedKeyType final {
              source_manager_.get_text_range(line_leading_whitespace),
              kAnnotation)});
   }
+
+  /**
+   * Adds the `@thrift.AllowUnsafeNonSealedKeyType` annotation before a
+   * function parameter.
+   *
+   * If the parameter is on its own line, delegates to `annotate_node`.
+   * If the parameter shares a line with the function signature (or other
+   * parameters), breaks the line after the preceding `(` or `,` and inserts
+   * the annotation on a new line before the parameter. The indentation is
+   * computed from the function's line (+ 2 spaces) so that all parameters
+   * align consistently.
+   */
+  void annotate_function_param(const t_field& param, const t_function& func) {
+    source_location param_loc = param.src_range().begin;
+    source_range line_ws = file_manager_.get_line_leading_whitespace(param_loc);
+
+    size_t param_offset = file_manager_.to_offset(param_loc);
+    size_t ws_end_offset = file_manager_.to_offset(line_ws.end);
+
+    // If param is on its own line, the existing logic works.
+    if (ws_end_offset >= param_offset) {
+      annotate_node(param_loc);
+      return;
+    }
+
+    maybe_add_include();
+
+    // Compute indent from the function's line so all params align at the
+    // same level, regardless of which line the param originally appeared on.
+    source_range func_line_ws =
+        file_manager_.get_line_leading_whitespace(func.src_range().begin);
+    std::string func_indent(source_manager_.get_text_range(func_line_ws));
+    std::string param_indent = func_indent + "  ";
+
+    // Scan backwards from param_loc over horizontal whitespace to find
+    // where the preceding '(' or ',' ends.
+    std::string_view old_content = file_manager_.old_content();
+    size_t scan_back = param_offset;
+    while (scan_back > 0 &&
+           (old_content[scan_back - 1] == ' ' ||
+            old_content[scan_back - 1] == '\t')) {
+      --scan_back;
+    }
+
+    // Replace whitespace between separator and param with
+    // line break + annotation + line break + indent.
+    file_manager_.add(
+        {.begin_pos = scan_back,
+         .end_pos = param_offset,
+         .new_content = fmt::format(
+             "\n{}{}\n{}", param_indent, kAnnotation, param_indent)});
+  }
 };
 
 } // namespace

third-party/thrift/src/thrift/compiler/codemod/annotate_allow_unsafe_non_sealed_key_type_test.py

@@ -325,3 +325,82 @@ def test_typedef_indirection_not_annotated(self):
                 }
                 """,
         )
+
+    def test_function_param_same_line(self):
+        self._check(
+            before="""\
+                include "thrift/annotation/thrift.thrift"
+
+                struct NonSealedStruct {}
+
+                service MyService {
+                  void myFunc(1: set<NonSealedStruct> param);
+                }
+                """,
+            after="""\
+                include "thrift/annotation/thrift.thrift"
+
+                struct NonSealedStruct {}
+
+                service MyService {
+                  void myFunc(
+                    @thrift.AllowUnsafeNonSealedKeyType
+                    1: set<NonSealedStruct> param);
+                }
+                """,
+        )
+
+    def test_function_param_same_line_multiple_params(self):
+        self._check(
+            before="""\
+                include "thrift/annotation/thrift.thrift"
+
+                struct NonSealedStruct {}
+
+                service MyService {
+                  void myFunc(1: i32 x, 2: set<NonSealedStruct> param);
+                }
+                """,
+            after="""\
+                include "thrift/annotation/thrift.thrift"
+
+                struct NonSealedStruct {}
+
+                service MyService {
+                  void myFunc(1: i32 x,
+                    @thrift.AllowUnsafeNonSealedKeyType
+                    2: set<NonSealedStruct> param);
+                }
+                """,
+        )
+
+    def test_function_param_same_line_multiple_non_sealed_params(self):
+        self._check(
+            before="""\
+                include "thrift/annotation/thrift.thrift"
+
+                struct NonSealedStruct {}
+
+                service MyService {
+                  void myFunc(1: set<NonSealedStruct> p1,
+                    2: set<NonSealedStruct> p2, 3: map<NonSealedStruct, i32> p3
+                  );
+                }
+                """,
+            after="""\
+                include "thrift/annotation/thrift.thrift"
+
+                struct NonSealedStruct {}
+
+                service MyService {
+                  void myFunc(
+                    @thrift.AllowUnsafeNonSealedKeyType
+                    1: set<NonSealedStruct> p1,
+                    @thrift.AllowUnsafeNonSealedKeyType
+                    2: set<NonSealedStruct> p2,
+                    @thrift.AllowUnsafeNonSealedKeyType
+                    3: map<NonSealedStruct, i32> p3
+                  );
+                }
+                """,
+        )