From 6b2b899874c2cf28beae851a963a8aa22bab8733 Mon Sep 17 00:00:00 2001
From: Chris Leary <cdleary@openai.com>
Date: Wed, 6 Nov 2024 19:19:14 -0800
Subject: [PATCH 1/2] [DSLX:TS] Prevent fatal when FunctionRef needs deduction
 and is not a function type.

---
 xls/dslx/type_system/deduce.cc                | 14 ++++++++++++--
 xls/dslx/type_system/type.cc                  |  4 ++++
 xls/dslx/type_system/type.h                   |  1 +
 xls/dslx/type_system/typecheck_module_test.cc |  8 ++++++++
 4 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/xls/dslx/type_system/deduce.cc b/xls/dslx/type_system/deduce.cc
index 39baf00234..da6581f1a5 100644
--- a/xls/dslx/type_system/deduce.cc
+++ b/xls/dslx/type_system/deduce.cc
@@ -2072,14 +2072,24 @@ class DeduceVisitor : public AstNodeVisitor {
     return Fatal(n);
   }
   absl::Status HandleFunctionRef(const FunctionRef* n) override {
-    return Fatal(n);
+    XLS_ASSIGN_OR_RETURN(std::unique_ptr<Type> callee_type,
+                         ctx_->Deduce(n->callee()));
+    if (!callee_type->IsFunction()) {
+      return TypeInferenceErrorStatus(
+          n->span(), callee_type.get(),
+          "Callee for function reference must be function-typed",
+          ctx_->file_table());
+    }
+    result_ = std::move(callee_type);
+    return absl::OkStatus();
   }
 
   absl::StatusOr<std::unique_ptr<Type>>& result() { return result_; }
 
  private:
   absl::Status Fatal(const AstNode* n) {
-    LOG(FATAL) << "Got unhandled AST node for deduction: " << n->ToString();
+    LOG(FATAL) << "DeduceVisitor got unhandled AST node for deduction: "
+               << n->ToString() << " node type name: " << n->GetNodeTypeName();
   }
 
   DeduceCtx* ctx_;
diff --git a/xls/dslx/type_system/type.cc b/xls/dslx/type_system/type.cc
index c4fecb9e52..7d3daec657 100644
--- a/xls/dslx/type_system/type.cc
+++ b/xls/dslx/type_system/type.cc
@@ -431,6 +431,10 @@ bool Type::IsTuple() const {
   return dynamic_cast<const TupleType*>(this) != nullptr;
 }
 
+bool Type::IsFunction() const {
+  return dynamic_cast<const FunctionType*>(this) != nullptr;
+}
+
 const EnumType& Type::AsEnum() const {
   auto* s = dynamic_cast<const EnumType*>(this);
   CHECK(s != nullptr) << "Type is not an enum: " << *this;
diff --git a/xls/dslx/type_system/type.h b/xls/dslx/type_system/type.h
index 79a21473e9..01f183945a 100644
--- a/xls/dslx/type_system/type.h
+++ b/xls/dslx/type_system/type.h
@@ -398,6 +398,7 @@ class Type {
   bool IsArray() const;
   bool IsMeta() const;
   bool IsTuple() const;
+  bool IsFunction() const;
 
   const StructType& AsStruct() const;
   const ProcType& AsProc() const;
diff --git a/xls/dslx/type_system/typecheck_module_test.cc b/xls/dslx/type_system/typecheck_module_test.cc
index a38b3de5fa..88aece7a54 100644
--- a/xls/dslx/type_system/typecheck_module_test.cc
+++ b/xls/dslx/type_system/typecheck_module_test.cc
@@ -3857,6 +3857,14 @@ proc t {
 )"));
 }
 
+TEST(TypecheckErrorTest, InvokingALiteralNumber) {
+  constexpr std::string_view kProgram = "fn f(x:u2) { 0<> }";
+  EXPECT_THAT(Typecheck(kProgram).status(),
+              IsPosError("TypeInferenceError",
+                         HasSubstr("Could not infer a type for this number, "
+                                   "please annotate a type.")));
+}
+
 TEST(TypecheckErrorTest, SignedValueToBuiltinExpectingUNViaParametric) {
   EXPECT_THAT(Typecheck(R"(
 fn p<S: bool, N: u32>() -> u32 {

From ab678d34a3bd28f59246d18ce4d09da39e73d766 Mon Sep 17 00:00:00 2001
From: Chris Leary <cdleary@openai.com>
Date: Wed, 6 Nov 2024 19:33:43 -0800
Subject: [PATCH 2/2] It is instantiation not invocation.

---
 xls/dslx/type_system/typecheck_module_test.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/xls/dslx/type_system/typecheck_module_test.cc b/xls/dslx/type_system/typecheck_module_test.cc
index 88aece7a54..2fe3c9e374 100644
--- a/xls/dslx/type_system/typecheck_module_test.cc
+++ b/xls/dslx/type_system/typecheck_module_test.cc
@@ -3857,7 +3857,7 @@ proc t {
 )"));
 }
 
-TEST(TypecheckErrorTest, InvokingALiteralNumber) {
+TEST(TypecheckErrorTest, InstantiateALiteralNumber) {
   constexpr std::string_view kProgram = "fn f(x:u2) { 0<> }";
   EXPECT_THAT(Typecheck(kProgram).status(),
               IsPosError("TypeInferenceError",