// Round-trip parsing/printing test.

// RUN: %target-sil-opt -sil-print-types %s -emit-sorted-sil | %FileCheck %s --check-prefix=CHECK-SIL

// Round-trip serialization-deserialization test.

// RUN: %empty-directory(%t)
// RUN: %target-sil-opt -sil-print-types %s -emit-sib -o %t/tmp.sib -module-name main
// RUN: %target-sil-opt -sil-print-types %t/tmp.sib -o %t/tmp.sil -module-name main

// https://github.com/apple/swift/issues/54526
// Workaround because import declarations are not preserved in .sib files.
// RUN: sed -e 's/import Swift$/import Swift; import _Differentiation/' %t/tmp.sil > %t/tmp_fixed.sil
// RUN: %target-sil-opt -sil-print-types %t/tmp_fixed.sil -module-name main -emit-sorted-sil | %FileCheck %s --check-prefix=CHECK-SIL

// IRGen test.

// RUN: %target-swift-frontend -emit-ir %s | %FileCheck %s --check-prefix=CHECK-IRGEN

// `shell` is required only to run `sed` as a
// https://github.com/apple/swift/issues/54526 workaround.
// REQUIRES: shell

// Would need to update for arm64e.
// UNSUPPORTED: CPU=arm64e

sil_stage raw

import Swift
import Builtin

import _Differentiation

sil @function : $@convention(thin) (Float) -> Float
sil @function_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)

sil @foo : $@convention(thin) (Float, Float, Float) -> (Float, Float)
sil @foo_vjp : $@convention(thin) (Float, Float, Float) -> (Float, Float, @owned @callee_guaranteed (Float) -> Float)

sil @make_differentiable_func : $@convention(thin) () -> @differentiable(reverse) @convention(thin) (Float, Float, @noDerivative Float) -> (Float, @noDerivative Float) {
// sil @make_differentiable_func : $@convention(thin) () -> () {
bb0:
  %orig_fn = function_ref @foo : $@convention(thin) (Float, Float, Float) -> (Float, Float)
  %vjp_fn = function_ref @foo_vjp : $@convention(thin) (Float, Float, Float) -> (Float, Float, @owned @callee_guaranteed (Float) -> Float)
  %diff_fn = differentiable_function [parameters 0 1] [results 0] %orig_fn : $@convention(thin) (Float, Float, Float) -> (Float, Float) with_derivative {undef : $@convention(thin) (Float, Float, Float) -> (Float, Float, @owned @callee_guaranteed (Float) -> Float), %vjp_fn : $@convention(thin) (Float, Float, Float) -> (Float, Float, @owned @callee_guaranteed (Float) -> Float)}
  %extracted_vjp = differentiable_function_extract [vjp] %diff_fn : $@differentiable(reverse) @convention(thin) (Float, Float, @noDerivative Float) -> (Float, @noDerivative Float)
  return %diff_fn : $@differentiable(reverse) @convention(thin) (Float, Float, @noDerivative Float) -> (Float, @noDerivative Float)
}

sil @make_differentiable_function : $@convention(thin) () -> @differentiable(reverse) @convention(thin) (Float) -> Float {
bb0:
  %orig_fn = function_ref @function : $@convention(thin) (Float) -> Float
  %vjp_fn = function_ref @function_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
  %diff_fn = differentiable_function [parameters 0] [results 0] %orig_fn : $@convention(thin) (Float) -> Float with_derivative {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), %vjp_fn : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
  %extracted_vjp = differentiable_function_extract [vjp] %diff_fn : $@differentiable(reverse) @convention(thin) (Float) -> Float
  %extracted_original = differentiable_function_extract [original] %diff_fn : $@differentiable(reverse) @convention(thin) (Float) -> Float
  return %diff_fn : $@differentiable(reverse) @convention(thin) (Float) -> Float
}

// CHECK-SIL-LABEL: @make_differentiable_function : $@convention(thin) () -> @differentiable(reverse) @convention(thin) (Float) -> Float {
// CHECK-SIL:   [[ORIG_FN:%.*]] = function_ref @function : $@convention(thin) (Float) -> Float
// CHECK-SIL:   [[VJP_FN:%.*]] = function_ref @function_vjp : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
// CHECK-SIL:   [[DIFF_FN:%.*]] = differentiable_function [parameters 0] [results 0] [[ORIG_FN]] : $@convention(thin) (Float) -> Float with_derivative {undef : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), [[VJP_FN]] : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}
// CHECK-SIL:   [[EXTRACTED_VJP_FN:%.*]] = differentiable_function_extract [vjp] [[DIFF_FN]] : $@differentiable(reverse) @convention(thin) (Float) -> Float
// CHECK-SIL:   [[EXTRACTED_ORIG_FN:%.*]] = differentiable_function_extract [original] [[DIFF_FN]] : $@differentiable(reverse) @convention(thin) (Float) -> Float
// CHECK-SIL:   return [[DIFF_FN]] : $@differentiable(reverse) @convention(thin) (Float) -> Float

// CHECK-IRGEN-LABEL: define {{.*}}swiftcc { ptr, ptr, ptr } @make_differentiable_function()
// CHECK-IRGEN-NEXT: entry:
// CHECK-IRGEN-NEXT:   ret { ptr, ptr, ptr } { ptr @function, ptr undef, ptr @function_vjp }

sil @examplefunc : $@convention(thin) (Float, Float, Float) -> Float
sil @examplemethod : $@convention(method) (Float, Float, Float) -> Float

// CHECK-SIL-LABEL: sil @test_roundtrip_parse
sil @test_roundtrip_parse : $@convention(thin) () -> () {
bb0:
  %0 = function_ref @examplefunc : $@convention(thin) (Float, Float, Float) -> Float
  %1 = differentiable_function [parameters 0 1 2] [results 0] %0 : $@convention(thin) (Float, Float, Float) -> Float with_derivative {undef : $@convention(thin) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float, Float, Float) -> Float), undef : $@convention(thin) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> (Float, Float, Float))}

  // CHECK-SIL: %2 = differentiable_function_extract [vjp] %1 : $@differentiable(reverse) @convention(thin) (Float, Float, Float) -> Float
  %2 = differentiable_function_extract [vjp] %1 : $@differentiable(reverse) @convention(thin) (Float, Float, Float) -> Float
  %3 = differentiable_function [parameters 0] [results 0] %0 : $@convention(thin) (Float, Float, Float) -> Float with_derivative {undef : $@convention(thin) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), undef : $@convention(thin) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}

  // CHECK-SIL: %4 = differentiable_function_extract [vjp] %3 : $@differentiable(reverse) @convention(thin) (Float, @noDerivative Float, @noDerivative Float) -> Float
  %4 = differentiable_function_extract [vjp] %3 : $@differentiable(reverse) @convention(thin) (Float, @noDerivative Float, @noDerivative Float) -> Float
  %5 = function_ref @examplemethod : $@convention(method) (Float, Float, Float) -> Float
  %6 = differentiable_function [parameters 0 1 2] [results 0] %5 : $@convention(method) (Float, Float, Float) -> Float with_derivative {undef : $@convention(method) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float, Float, Float) -> Float), undef : $@convention(method) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> (Float, Float, Float))}

  // CHECK-SIL: %7 = differentiable_function_extract [vjp] %6 : $@differentiable(reverse) @convention(method) (Float, Float, Float) -> Float
  %7 = differentiable_function_extract [vjp] %6 : $@differentiable(reverse) @convention(method) (Float, Float, Float) -> Float
  %8 = differentiable_function [parameters 0] [results 0] %5 : $@convention(method) (Float, Float, Float) -> Float with_derivative {undef : $@convention(method) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), undef : $@convention(method) (Float, Float, Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)}

  // CHECK-SIL: %9 = differentiable_function_extract [vjp] %8 : $@differentiable(reverse) @convention(method) (Float, @noDerivative Float, @noDerivative Float) -> Float
  %9 = differentiable_function_extract [vjp] %8 : $@differentiable(reverse) @convention(method) (Float, @noDerivative Float, @noDerivative Float) -> Float

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