# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -o - | FileCheck %s

node {
  name: "Call_foo"
  op: "PartitionedCall"
  attr {
    key: "Tin"
    value {
      list {
      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_INT32
      }
    }
  }
  attr {
    key: "f"
    value {
      func {
        name: "foo"
      }
    }

  func.func @main() {
    tf_executor.graph {
      %0 = tf_executor.island wraps "tf.PartitionedCall"() {Tin = [], Tout = [], config = "", config_proto = "", device = "", executor_type = "", f = @foo, name = "Call_foo"} : () -> ()
      tf_executor.fetch
    }
    func.return
  }
  func.func @foo() {
    tf_executor.graph {
      %0:2 = tf_executor.island {

  // CHECK:          name: "foo"
  // CHECK:        }
    %1:2 = tf_executor.island wraps "tf.PartitionedCall"(%0) {Tin = [], Tout = [], config = "", config_proto = "", device = "", executor_type = "", f = @foo, name = "Call_foo"} : (tensor<f32>) -> tensor<*xf32>
    tf_executor.fetch
  }
  func.return
}

// CHECK:      library {
// CHECK-NEXT:   function {
// CHECK-NEXT:     signature {
// CHECK-NEXT:       name: "foo"

    tf_executor.fetch
  }
  func.return
}

func.func @foo(%arg0: tensor<!tf_type.resource>) -> tensor<!tf_type.resource> {
  %graph = tf_executor.graph {
    tf_executor.fetch %arg0 : tensor<!tf_type.resource>

    // CHECK:     name: "foo"
    // CHECK:   }
    // CHECK: }
    %0 = tf_executor.island wraps "tf.PartitionedCall"() {Tin = [], Tout = [], config = "", config_proto = "", device = "", executor_type = "", f = @foo, name = "Call_foo"} : () -> ()
    tf_executor.fetch
  }
  func.return
}

// CHECK:      library {
// CHECK:        function {
// CHECK-NEXT:     signature {
// CHECK-NEXT:       name: "bar"
// CHECK-NEXT:     }

            <caret_context>val x = 0
        }
    }
}

fun callFoo(foo: Foo): Int {
    return 0
}

fun callString(string: String): Int {
    return 1
}


// MODULE: main
// MODULE_KIND: CodeFragment
// CONTEXT_MODULE: context

// FILE: fragment.kt
// CODE_FRAGMENT_KIND: EXPRESSION

              $this: CALL 'public final fun plus (other: kotlin.Int): kotlin.Int [operator] declared in kotlin.Int' type=kotlin.Int origin=PLUS
                $this: CALL 'public final fun callFoo (foo: <root>.Foo): kotlin.Int declared in <root>.ContextKt' type=kotlin.Int origin=null
                  foo: GET_VAR 'p0: <root>.Foo declared in <root>.CodeFragment.run' type=<root>.Foo origin=null