[AutoDiff] Fix assert on missing struct decl on cross-file derivative search (#77183)

Consider:

1. File struct.swift defining `struct Struct` with `static func max` member
2. File derivatives.swift defining `extension Struct` with custom derivative of the `max` function
3. File error.swift defining a differentiable function which uses `Struct.max`.

Previously, when passing error.swift as primary file and derivatives.swift as a secondary file to swift-frontend (and forgetting to pass struct.swift as a secondary file as well), an assertion failure was triggered.

This patch fixes the issue by adding a check against `ErrorType` in `findAutoDiffOriginalFunctionDecl` before calling `lookupMember`.

Co-authored-by: Anton Korobeynikov <[email protected]>

Author: kovdan01

lib/Sema/TypeCheckAttr.cpp

@@ -5460,7 +5460,8 @@ static AbstractFunctionDecl *findAutoDiffOriginalFunctionDecl(
   if (!baseType && lookupContext->isTypeContext())
     baseType = lookupContext->getSelfTypeInContext();
   if (baseType) {
-    results = TypeChecker::lookupMember(lookupContext, baseType, funcName);
+    if (!baseType->hasError())
+      results = TypeChecker::lookupMember(lookupContext, baseType, funcName);
   } else {
     results = TypeChecker::lookupUnqualified(
         lookupContext, funcName, funcNameLoc.getBaseNameLoc(), lookupOptions);

test/AutoDiff/Sema/DerivativeRegistrationCrossFile/Inputs/derivatives-error.swift

@@ -0,0 +1,37 @@
+import _Differentiation
+
+@inlinable
+@derivative(of: min)
+func minVJP<T: Comparable & Differentiable>(
+    _ x: T,
+    _ y: T
+) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
+    func pullback(_ v: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
+        if x <= y {
+            return (v, .zero)
+        }
+        else {
+            return (.zero, v)
+        }
+    }
+    return (value: min(x, y), pullback: pullback)
+}
+
+extension Struct {
+    @inlinable
+    @derivative(of: max) // expected-error {{cannot find 'max' in scope}}
+    static func maxVJP<T: Comparable & Differentiable>(
+        _ x: T,
+        _ y: T
+    ) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
+        func pullback(_ v: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
+            if x < y {
+                return (.zero, v)
+            }
+            else {
+                return (v, .zero)
+            }
+        }
+        return (value: max(x, y), pullback: pullback)
+    }
+}

test/AutoDiff/Sema/DerivativeRegistrationCrossFile/Inputs/derivatives.swift

@@ -17,19 +17,21 @@ func minVJP<T: Comparable & Differentiable>(
     return (value: min(x, y), pullback: pullback)
 }
 
-@inlinable
-@derivative(of: max)
-func maxVJP<T: Comparable & Differentiable>(
-    _ x: T,
-    _ y: T
-) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
-    func pullback(_ v: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
-        if x < y {
-            return (.zero, v)
-        }
-        else {
-            return (v, .zero)
+extension Struct {
+    @inlinable
+    @derivative(of: max)
+    static func maxVJP<T: Comparable & Differentiable>(
+        _ x: T,
+        _ y: T
+    ) -> (value: T, pullback: (T.TangentVector) -> (T.TangentVector, T.TangentVector)) {
+        func pullback(_ v: T.TangentVector) -> (T.TangentVector, T.TangentVector) {
+            if x < y {
+                return (.zero, v)
+            }
+            else {
+                return (v, .zero)
+            }
         }
+        return (value: max(x, y), pullback: pullback)
     }
-    return (value: max(x, y), pullback: pullback)
 }

test/AutoDiff/Sema/DerivativeRegistrationCrossFile/Inputs/struct.swift

@@ -0,0 +1,14 @@
+import _Differentiation
+
+struct Struct {
+    @inlinable
+    static func max<T: Comparable>(
+        _ x: T,
+        _ y: T
+    ) -> T {
+        if x > y
+            return y
+        else
+            return x
+    }
+}

test/AutoDiff/Sema/DerivativeRegistrationCrossFile/error.swift

@@ -0,0 +1,8 @@
+// RUN: %target-swift-frontend -emit-sil -verify -primary-file %s %S/Inputs/derivatives-error.swift -module-name main -o /dev/null
+
+import _Differentiation
+
+@differentiable(reverse)
+func clamp(_ value: Double, _ lowerBound: Double, _ upperBound: Double) -> Double {
+    return Struct.max(min(value, upperBound), lowerBound) // expected-error {{cannot find 'Struct' in scope}}
+}

test/AutoDiff/Sema/DerivativeRegistrationCrossFile/main.swift

@@ -1,9 +1,10 @@
-// RUN: %target-swift-frontend -emit-sil -verify -primary-file %s %S/Inputs/derivatives.swift -module-name main -o /dev/null
+// RUN: %target-swift-frontend -emit-sil -verify -primary-file %s \
+// RUN:   %S/Inputs/derivatives.swift %S/Inputs/struct.swift -module-name main -o /dev/null
 
 import _Differentiation
 
 @differentiable(reverse)
 func clamp(_ value: Double, _ lowerBound: Double, _ upperBound: Double) -> Double {
     // No error expected
-    return max(min(value, upperBound), lowerBound)
+    return Struct.max(min(value, upperBound), lowerBound)
 }