// RUN: %target-sil-opt -differentiation -differentiation-skip-folding-differentiable-function-extraction %s | %FileCheck %s

sil_stage raw

import Builtin
import Swift
import SwiftShims

protocol DiffReq {
  @differentiable(wrt: (x))
  func f(_ x: Float) -> Float
}

sil @differentiateWitnessMethod : $@convention(thin) <T where T : DiffReq> (@in_guaranteed T) -> () {
bb0(%0 : $*T):
  %1 = witness_method $T, #DiffReq.f : <Self where Self : DiffReq> (Self) -> (Float) -> Float : $@convention(witness_method: DiffReq) <τ_0_0 where τ_0_0 : DiffReq> (Float, @in_guaranteed τ_0_0) -> Float
  %2 = differentiable_function [parameters 0] [results 0] %1 : $@convention(witness_method: DiffReq) <τ_0_0 where τ_0_0 : DiffReq> (Float, @in_guaranteed τ_0_0) -> Float

  %ret = tuple ()
  return %ret : $()
}

// CHECK-LABEL: sil @differentiateWitnessMethod
// CHECK:   [[ORIG_REF:%.*]] = witness_method $T, #DiffReq.f
// CHECK:   [[JVP_REF:%.*]] = witness_method $T, #DiffReq.f!jvp.SU
// CHECK:   [[VJP_REF:%.*]] = witness_method $T, #DiffReq.f!vjp.SU
// CHECK:   differentiable_function [parameters 0] [results 0] [[ORIG_REF]] : {{.*}} with_derivative {[[JVP_REF]] : {{.*}}, [[VJP_REF]] : {{.*}}}
// CHECK: } // end sil function 'differentiateWitnessMethod'

sil @differentiatePartiallyAppliedWitnessMethod : $@convention(thin) <T where T : DiffReq> (@in_guaranteed T) -> () {
bb0(%0 : $*T):
  %1 = witness_method $T, #DiffReq.f : <Self where Self : DiffReq> (Self) -> (Float) -> Float : $@convention(witness_method: DiffReq) <τ_0_0 where τ_0_0 : DiffReq> (Float, @in_guaranteed τ_0_0) -> Float
  %2 = partial_apply [callee_guaranteed] %1<T>(%0) : $@convention(witness_method: DiffReq) <τ_0_0 where τ_0_0 : DiffReq> (Float, @in_guaranteed τ_0_0) -> Float
  %3 = differentiable_function [parameters 0] [results 0] %2 : $@callee_guaranteed (Float) -> Float

  %ret = tuple ()
  return %ret : $()
}

// CHECK-LABEL: sil @differentiatePartiallyAppliedWitnessMethod
// CHECK: bb0([[ARG:%.*]] : $*T):
// CHECK:  [[ORIG_REF:%.*]] = witness_method $T, #DiffReq.f
// CHECK:  [[ARGCOPY1:%.*]] = alloc_stack $T
// CHECK:  copy_addr [[ARG]] to [initialization] [[ARGCOPY1]] : $*T
// CHECK:  [[ARGCOPY2:%.*]] = alloc_stack $T
// CHECK:  copy_addr [[ARG]] to [initialization] [[ARGCOPY2]] : $*T
// CHECK:  [[ORIG_REF_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[ORIG_REF]]<T>(%0)
// CHECK:  [[JVP_REF:%.*]] = witness_method $T, #DiffReq.f!jvp.SU
// CHECK:  [[JVP_REF_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[JVP_REF]]<T>([[ARGCOPY1]])
// CHECK:  [[VJP_REF:%.*]] = witness_method $T, #DiffReq.f!vjp.SU
// CHECK:  [[VJP_REF_PARTIALLY_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[VJP_REF]]<T>([[ARGCOPY2]])
// CHECK:  dealloc_stack [[ARGCOPY2]]
// CHECK:  dealloc_stack [[ARGCOPY1]]
// CHECK:  differentiable_function [parameters 0] [results 0] [[ORIG_REF_PARTIALLY_APPLIED]] : {{.*}} with_derivative {[[JVP_REF_PARTIALLY_APPLIED]] : {{.*}}, [[VJP_REF_PARTIALLY_APPLIED]] : {{.*}}}
// CHECK: } // end sil function 'differentiatePartiallyAppliedWitnessMethod'