[AutoDiff] Runtime crash in certain cases when attempting a Dictionary lookup.

Description

The code below crashes during runtime. A second, distinct runtime crash can be generated by editing one line in the reproducer below.

Reproduction

Compile the code below in Debug mode and attempt to run it.

import _Differentiation; import Foundation;
var z: any E = F();
print(try y(b: K.self, i: &z))
func x<N: G>(p _: N.Type, M: inout N.L.M) throws -> [D<Double, Double>]
{
    let a = D<Double, Double>()
    @differentiable(reverse, wrt: (p)) func l(p: N.L.M,  t: D<Double, Double>) -> D<Double, Double> {var r = t; s(f: p, t: &r); return r}
    func q(q: N.L.M) -> D<Double, Double> {return l(p: q, t: a)}
    print(valueWithPullback(at: M, of: q))
    return [D<Double, Double>]()
}
@differentiable(reverse, wrt: (M, d)) func s<M>(f M: M, t d: inout D<Double, Double>) where M: J {
    let b = B<F, Double>(k: [\F.f])
    d.p((b as! B<M, Double>).r(f: M), i: 0);
    //d.p(3.3, i: 0);
}
func y(b: E.Type, i: inout E) throws -> [D<Double, Double>] {return try (b as! any I.Type).d().c(i: &i)}
func h<O, M>(_ o: O, at m: WritableKeyPath<O, M>) -> M {return o[keyPath: m]}
@derivative(of: h)  func w<O, M>(_ o: O, at m: WritableKeyPath<O, M>) -> (value: M, pullback: (M.TangentVector) -> O.TangentVector) where O: J, M: Differentiable
{
    var z = o.z
    let c = Q[m] as! WritableKeyPath<O.TangentVector, M.TangentVector>
    return (value: o[keyPath: m], pullback: { d in z[keyPath: c] = d; return z})
}
var Q: [AnyKeyPath: AnyKeyPath] = [:]
struct A<L>: G where L: I {}
struct B<R, V: E>{@noDerivative var k: [WritableKeyPath<R, V>]}
struct C<V> { var w: V}
struct D<I: Hashable, DataType> {}
protocol E{}
struct F : J & E {var f: Double = 2}
protocol G where L: I {associatedtype L: H}
protocol H {associatedtype M: E}
protocol I: H where M: J {}
protocol J: Differentiable {var u: () -> TangentVector { get }}
extension D: Differentiable where DataType: Differentiable & J, DataType == DataType.TangentVector {@differentiable(reverse) mutating func p(_ v: DataType, i: I) {}}
extension J {@inlinable var u: () -> TangentVector {{ Self.TangentVector.zero }}; @inlinable var z: TangentVector { u() }}
enum K: E & I {typealias M = F} //Signal 5 if conformance to I is removed.
extension Double: E & J {}
extension C: Differentiable where V: Differentiable {
    typealias TangentVector = V.TangentVector
    mutating func move(by o: V.TangentVector) {}
    @differentiable(reverse) init(a: V) {self.init(w: a)}
    @derivative(of: init(a:)) static func v(a: V) -> (value: Self, pullback: (Self.TangentVector) -> V.TangentVector) {(value: .init(a: a), pullback: { $0 })}
    @differentiable(reverse) var d: V {self.w}
    @derivative(of: d) func f() -> (value: V, pullback: (V.TangentVector) -> Self.TangentVector) {(value: self.w, pullback: { $0 })}
}
extension B where R: J, V: Differentiable {
    @differentiable(reverse) func r(f r: R) -> V {C<V>(a: self.s(f: r)[0]).d}
    @differentiable(reverse) func s(f r: R) -> [V] {var s: [V] = []; for k in self.k {let t: V = h(r, at: k); s = s + [t]}; return s}
}
extension I {static func d() -> any G.Type {return A<Self>.self}}
extension G {static func c(i: inout E) throws -> [D<Double, Double>] {var M = i as! L.M; return try x(p: Self.self, M: &M)}}

Stack dump

I have two stack traces, both from a Ubuntu 20 environment. These were generated via lldb. The first is from the above reproducer as-is:

    * thread #1, name = 'main', stop reason = signal SIGSEGV: address not mapped to object (fault address: 0x10)
  * frame #0: 0x00007ffff7a6b1a1 libswiftCore.so`Swift._NativeDictionary.lookup(τ_0_0) -> Swift.Optional<τ_0_1> + 17
    frame #1: 0x0000555555559bbf main`w<O, M>(o=(f = 2), m=0x0000555555585de0) at main.swift:22:14
    frame #2: 0x000055555555a019 main`reverse-mode derivative of h<A, B>(_:at:) at <compiler-generated>:0
    frame #3: 0x000055555555d9c7 main`reverse-mode derivative of B<>.s(r=(f = 2), self=main.B<main.F, Swift.Double> @ 0x00007fffffffd968) at main.swift:50:98
    frame #4: 0x000055555555cc10 main`partial apply at <compiler-generated>:0
    frame #5: 0x000055555555c591 main`reverse-mode derivative of B<>.r(r=(f = 2), self=main.B<main.F, Swift.Double> @ 0x00007fffffffdbc8) at main.swift:49:64
    frame #6: 0x000055555555bdb9 main`partial apply at <compiler-generated>:0
    frame #7: 0x000055555555bb13 main`reverse-mode derivative of s<τ_0_0>(M=(f = 2), d=main.D<Swift.Double, Swift.Double> @ scalar) at main.swift:14:30
    frame #8: 0x000055555555b3dc main`reverse-mode derivative of l #1 <τ_0_0>(p=(f = 2), t=main.D<Swift.Double, Swift.Double> @ scalar) in x<τ_0_0>(p:M:) at main.swift:7:113
    frame #9: 0x000055555555f267 main`reverse-mode derivative of q #1 <τ_0_0>(q=(f = 2), a=main.D<Swift.Double, Swift.Double> @ scalar) in x<τ_0_0>(p:M:) at main.swift:8:51
    frame #10: 0x00005555555590da main`thunk for @callee_guaranteed (@in_guaranteed A.G.L.H.M) -> (@unowned D<Double, Double>, @owned @escaping @callee_guaranteed @substituted <A> (@unowned D<Double, Double><>.TangentVector) -> (@out A) for <A.G.L.H.M.Differentiable.TangentVector>) at <compiler-generated>:0
    frame #11: 0x00007ffff7f840e0 libswift_Differentiation.so`_Differentiation.valueWithPullback<τ_0_0, τ_0_1 where τ_0_0: _Differentiation.Differentiable, τ_0_1: _Differentiation.Differentiable>(at: τ_0_0, of: @differentiable(reverse) (τ_0_0) -> τ_0_1) -> (value: τ_0_1, pullback: (τ_0_1.TangentVector) -> τ_0_0.TangentVector) + 96
    frame #12: 0x0000555555558aa4 main`x<N>(_0=@thick main.A<main.K>.Type, M=(f = 2)) at main.swift:9:11
    frame #13: 0x00005555555594f2 main`static G.c(i=(payload_data_0 = 0x4000000000000000, payload_data_1 = 0x0000000000000000, payload_data_2 = 0x0000000000000000, metadata = Any.Type, wtable = 0x000055555556f958 main`protocol witness table for main.F : main.E in main), self=@thick main.A<main.K>.Type) at main.swift:53:101
    frame #14: 0x0000555555558553 main`y(b=main.K, i=(payload_data_0 = 0x4000000000000000, payload_data_1 = 0x0000000000000000, payload_data_2 = 0x0000000000000000, metadata = Any.Type, wtable = 0x000055555556f958 main`protocol witness table for main.F : main.E in main)) at main.swift:17:96
    frame #15: 0x0000555555558386 main`main at main.swift:3:11
    frame #16: 0x00007ffff69c2083 libc.so.6`__libc_start_main + 243
    frame #17: 0x000055555555803e main`_start + 46

The second is the result of removing conformance to I in line 38:

* thread #1, name = 'main', stop reason = Swift runtime failure: type cast failed
    frame #0: 0x000055555555862f main`dynamic_cast_existential_1_unconditional [inlined] Swift runtime failure: type cast failed at <compiler-generated>:0
  * frame #1: 0x000055555555862f main`dynamic_cast_existential_1_unconditional at <compiler-generated>:0
    frame #2: 0x00005555555584a2 main`y(b=main.K, i=(payload_data_0 = 0x4000000000000000, payload_data_1 = 0x0000000000000000, payload_data_2 = 0x0000000000000000, metadata = Any.Type, wtable = 0x000055555556f958 main`protocol witness table for main.F : main.E in main)) at main.swift:17:76
    frame #3: 0x00005555555582f6 main`main at main.swift:3:11
    frame #4: 0x00007ffff69c2083 libc.so.6`__libc_start_main + 243
    frame #5: 0x0000555555557fae main`_start + 46

Expected behavior

The program should print at least

(value: projectName.D<Swift.Double, Swift.Double>(), pullback: (Function))
[]

and exit with an exit code of 0. This is the result of uncommenting line 15, and commenting line 14. (Remember to restore the conformance to I in line 38 if removed).

Environment

This affects at least macOS and Ubuntu (20). Toolchains from 2023-12-07a (and possibly earlier) to 2024-06-08 are affected.

Additional information

A multi-package version of this reproducer was generating a runtime segfault when attempting to dereference pointer 0x0 rather than 0x10 when attempting to perform the Dictionary lookup.

swiftlang / swift