Properly clone and format the extern Verilog attribute (Foreign Funct…

…ion Interface/FFI) of functions.

Fixes #1741

PiperOrigin-RevId: 699269626

xls/dslx/fmt/ast_fmt.cc

@@ -2007,23 +2007,33 @@ DocRef Formatter::Format(const Function& n) {
         arena_.MakeNest(ConcatNGroup(arena_, params_pieces)));
   }
 
-  // For empty function we don't put spaces between the curls.
-  if (n.body()->empty()) {
-    std::vector<DocRef> fn_pieces = {
-        ConcatNGroup(arena_, signature_pieces),
-        FmtBlock(*n.body(), comments_, arena_, /*add_curls=*/false),
-        arena_.ccurl(),
-    };
-
-    return ConcatNGroup(arena_, fn_pieces);
+  std::vector<DocRef> fn_pieces;
+  if (n.extern_verilog_module().has_value()) {
+    auto code_template = (*n.extern_verilog_module()).code_template();
+    fn_pieces.push_back(
+        ConcatN(arena_, {
+                            arena_.MakeText("#[extern_verilog(\""),
+                            arena_.MakeText(code_template),
+                            arena_.MakeText("\")]"),
+                            arena_.hard_line(),
+                        }));
   }
 
-  std::vector<DocRef> fn_pieces = {
-      ConcatNGroup(arena_, signature_pieces),
-      arena_.break1(),
-      FmtBlock(*n.body(), comments_, arena_, /*add_curls=*/false),
-      arena_.break1(),
-      arena_.ccurl(),
+  if (n.body()->empty()) {
+    // For empty function we don't put spaces between the curls.
+    fn_pieces.push_back(ConcatNGroup(arena_, signature_pieces));
+    fn_pieces.push_back(
+        FmtBlock(*n.body(), comments_, arena_, /*add_curls=*/false));
+    fn_pieces.push_back(arena_.ccurl());
+  } else {
+    // For non-empty functions, we break after the signature and before
+    // the ccurl.
+    fn_pieces.push_back(ConcatNGroup(arena_, signature_pieces));
+    fn_pieces.push_back(arena_.break1());
+    fn_pieces.push_back(
+        FmtBlock(*n.body(), comments_, arena_, /*add_curls=*/false));
+    fn_pieces.push_back(arena_.break1());
+    fn_pieces.push_back(arena_.ccurl());
   };
 
   return ConcatNGroup(arena_, fn_pieces);

xls/dslx/fmt/ast_fmt_test.cc

@@ -2754,5 +2754,19 @@ TEST_F(ModuleFmtTest, TooLongLines) {
 )");
 }
 
+TEST_F(ModuleFmtTest, ExternVerilog) {
+  // Note alternate string literal delimiter * so it can use )" on the
+  // last line of the annotation.
+  DoFmt(R"*(#[extern_verilog("external_divmod #(
+     .divisor_width({B_WIDTH})
+    ) {fn} (
+     .dividend({a}),
+     .by_zero({return.1})
+    )")]
+fn divmod() {
+}
+)*");
+}
+
 }  // namespace
 }  // namespace xls::dslx

xls/dslx/frontend/ast_cloner.cc

@@ -393,10 +393,14 @@ class AstCloner : public AstNodeVisitor {
       new_return_type =
           down_cast<TypeAnnotation*>(old_to_new_.at(n->return_type()));
     }
-    old_to_new_[n] = module_->Make<Function>(
+    auto new_function = module_->Make<Function>(
         n->span(), new_name_def, new_parametric_bindings, new_params,
         new_return_type, down_cast<StatementBlock*>(old_to_new_.at(n->body())),
         n->tag(), n->is_public());
+    if (n->extern_verilog_module().has_value()) {
+      new_function->set_extern_verilog_module(*n->extern_verilog_module());
+    }
+    old_to_new_[n] = new_function;
     new_name_def->set_definer(old_to_new_.at(n));
     return absl::OkStatus();
   }

xls/dslx/frontend/ast_cloner_test.cc

@@ -1636,5 +1636,26 @@ struct Point {
   EXPECT_EQ(kProgram, clone->ToString());
 }
 
+TEST(AstClonerTest, ExternVerilog) {
+  // Note alternate string literal delimiter * so it can use )" on the
+  // last line of the annotation.
+  constexpr std::string_view kProgram =
+      R"*(#[extern_verilog("external_divmod #(
+     .divisor_width({B_WIDTH})
+    ) {fn} (
+     .dividend({a}),
+     .by_zero({return.1})
+    )")]
+fn divmod() {})*";
+
+  FileTable file_table;
+  XLS_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<Module> module,
+      ParseModule(kProgram, "fake_path.x", "the_module", file_table));
+  XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr<Module> clone,
+                           CloneModule(*module.get()));
+  EXPECT_EQ(kProgram, clone->ToString());
+}
+
 }  // namespace
 }  // namespace xls::dslx

xls/examples/BUILD

@@ -319,6 +319,11 @@ xls_dslx_test(
     dslx_test_args = {"compare": "jit"},
 )
 
+xls_dslx_fmt_test(
+    name = "ffi_fmt_test",
+    src = "ffi.x",
+)
+
 xls_dslx_verilog(
     name = "ffi_codegen",
     codegen_args = {

xls/examples/ffi.x

@@ -15,7 +15,6 @@
 // Example how to instantiate an existing Verilog module, with the DSLX
 // foreign function interface.
 
-
 // A function that we want to interface with some existing verilog module
 // we implement the function itself in DSLX and provide a text template that
 // will be used to create an instantiation of the module with the correct
@@ -47,25 +46,18 @@
      .remainder({return.0.1}),  // ... values referenced as return.0 and return.1
      .by_zero({return.1})
     )")]
-fn divmod<A_WIDTH:u32, B_WIDTH:u32>(a:bits[A_WIDTH], b:bits[B_WIDTH])
-                                    -> ((bits[A_WIDTH], bits[B_WIDTH]), bool) {
-    if (b == u32:0) {
-        ((a, b), true)
-    } else {
-        ((a / b, a % b), false)
-    }
+fn divmod<A_WIDTH: u32, B_WIDTH: u32>
+    (a: bits[A_WIDTH], b: bits[B_WIDTH]) -> ((bits[A_WIDTH], bits[B_WIDTH]), bool) {
+    if b == u32:0 { ((a, b), true) } else { ((a / b, a % b), false) }
 }
 
-
-fn main(dividend:u32, divisor:u32) -> ((u32, u32), bool) {
-  divmod(dividend, divisor)
-}
+fn main(dividend: u32, divisor: u32) -> ((u32, u32), bool) { divmod(dividend, divisor) }
 
 // In regular testing and bytecode interpretation, the DSLX function is used.
 // Only in the finaly code-generation, the function invocation is replaced with
 // the user-chosen module.
 #[test]
 fn divmod_test() {
-  assert_eq(divmod(u32:42, u32:5), ((u32:8, u32:2), false));
-  assert_eq(divmod(u32:42, u32:0), ((u32:42, u32:0), true));
+    assert_eq(divmod(u32:42, u32:5), ((u32:8, u32:2), false));
+    assert_eq(divmod(u32:42, u32:0), ((u32:42, u32:0), true));
 }