Support returning Self type from impl functions.

google#617

PiperOrigin-RevId: 704879436

xls/dslx/type_system/deduce.cc

@@ -1798,16 +1798,24 @@ absl::StatusOr<std::unique_ptr<Type>> DeduceSelfTypeAnnotation(
     const SelfTypeAnnotation* node, DeduceCtx* ctx) {
   VLOG(5) << "DeduceSelfTypeAnnotation: " << node->ToString();
 
-  // Currently `self` is only supported for functions within `impl`. Identify
-  // the relevant `impl` and return the type for the associated `struct`.
-  // Lineage:
-  // impl --> function --> parameter --> type.
-  AstNode* param = node->parent();
-  XLS_RET_CHECK(param != nullptr);
-  AstNode* fn = param->parent();
+  // Currently `self` is only supported for functions within `impl` either as a
+  // parameter or a return type.
+  //  * Lineage for a parameter:
+  //    impl --> function --> parameter --> type.
+  //  * Lineage for a function return:
+  //    impl --> function --> type.
+  //
+  // Identify the relevant `impl` and return the type for the associated
+  // `struct`.
+
+  Function* fn;
+  if (auto* param = dynamic_cast<Param*>(node->parent())) {
+    fn = dynamic_cast<Function*>(param->parent());
+  } else {
+    fn = dynamic_cast<Function*>(node->parent());
+  }
   XLS_RET_CHECK(fn != nullptr);
-  AstNode* impl_node = fn->parent();
-  Impl* impl = dynamic_cast<Impl*>(impl_node);
+  Impl* impl = dynamic_cast<Impl*>(fn->parent());
   XLS_RET_CHECK(impl != nullptr);
 
   return ctx->Deduce(ToAstNode(impl->struct_ref()));

xls/dslx/type_system/impl_typecheck_test.cc

@@ -156,6 +156,22 @@ fn main() -> uN[8] {
       ParseAndTypecheck(kProgram, "fake_main_path.x", "main", &import_data));
 }
 
+TEST(TypecheckTest, ImplFunctionReturnsSelf) {
+  constexpr std::string_view kProgram = R"(
+struct Point { x: u32, y: u32 }
+
+impl Point {
+    fn unit() -> Self { Point { x: u32:1, y: u32:1 } }
+}
+
+fn point_dims() -> u1 {
+    let p = Point::unit();
+    uN[p.x]:0
+}
+)";
+  XLS_EXPECT_OK(Typecheck(kProgram));
+}
+
 TEST(TypecheckTest, ImplsForDifferentStructs) {
   constexpr std::string_view kProgram = R"(
 struct Point { x: u32, y: u32 }