// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// RUN: litert-opt --tfl-legalize-tf-while %s -o - | FileCheck %s
// RUN: litert-opt --tfl-legalize-tf-while %s -o - --tfl-legalize-tf-while --inline='default-pipeline=''' | FileCheck %s --check-prefix=INLINE
// RUN: litert-opt --tfl-legalize-tf-while %s -o - --tfl-legalize-tf-while --inline | FileCheck %s --check-prefix=CANON

func.func @while_main(%arg0: tensor<?x256x256xf32>) -> (tensor<i32>, tensor<256x256xf32>, tensor<?x256x256xf32>) attributes {tf.entry_function = {inputs = "input", outputs = "Identity,Identity_1,Identity_2"}} {
  %cst = arith.constant dense<1.000000e+00> : tensor<256x256xf32>
  %cst_0 = arith.constant dense<0> : tensor<i32>
  %cst_1 = arith.constant dense<-1> : tensor<i32>
  %0:5 = "tf.While"(%cst_0, %cst_1, %cst_0, %cst, %arg0) {body = @while_body_11_frozen0, cond = @while_cond_10_frozen0, device = "", is_stateless = true} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<256x256xf32>, tensor<?x256x256xf32>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<256x256xf32>, tensor<?x256x256xf32>)
  func.return %0#2, %0#3, %0#4 : tensor<i32>, tensor<256x256xf32>, tensor<?x256x256xf32>
}

func.func @while_body_11_frozen0(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<*xi32>, %arg3: tensor<*xf32>, %arg4: tensor<*xf32>) -> (tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*xf32>, tensor<*xf32>) {
  %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
  %cst_0 = arith.constant dense<0> : tensor<i32>
  %cst_1 = arith.constant dense<-1> : tensor<i32>
  %cst_2 = arith.constant dense<1> : tensor<i32>
  %0 = "tf.AddV2"(%arg2, %cst_2) {T = i32, device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
  %1 = "tf.Transpose"(%arg3, %cst) {T = f32, Tperm = i32, device = ""} : (tensor<*xf32>, tensor<2xi32>) -> tensor<?x?xf32>
  %2 = "tf.Rank"(%arg3) : (tensor<*xf32>) -> tensor<i32>
  %3 = "tf.Range"(%2, %cst_0, %cst_1) : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
  %4 = "tf.Sub"(%3, %cst_2) : (tensor<?xi32>, tensor<i32>) -> tensor<?xi32>
  %5 = "tf.Transpose"(%arg3, %4) : (tensor<*xf32>, tensor<?xi32>) -> tensor<*xf32>
  %6 = "tf.MatMul"(%1, %5) {transpose_a = false, transpose_b = true} : (tensor<?x?xf32>, tensor<*xf32>) -> tensor<?x?xf32>
  %7 = "tf.AddV2"(%arg4, %6) {T = f32, device = ""} : (tensor<*xf32>, tensor<?x?xf32>) -> tensor<*xf32>
  %8 = "tf.AddV2"(%arg0, %cst_2) {T = i32, device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
  func.return %8, %arg1, %0, %arg3, %7 : tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*xf32>, tensor<*xf32>
}

func.func @while_cond_10_frozen0(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<*xi32>, %arg3: tensor<*xf32>, %arg4: tensor<*xf32>) -> tensor<*xi1> {
  %cst = arith.constant dense<10> : tensor<i32>
  %0 = "tf.Less"(%arg2, %cst) {T = i32, device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
  func.return %0 : tensor<*xi1>
}

// CHECK: tfl.while
// CHECK: ^bb0([[ARGS:.*]]):
// CHECK:   call @while_cond_10_frozen0
// CHECK:   yield
// CHECK: ^bb0([[ARGS]]):
// CHECK:   call @while_body
// CHECK:   yield
// CHECK: while_body
// CHECK: while_cond

// INLINE: tfl.while
// INLINE: ^bb0([[ARGS:.*]]):
// INLINE:   tf.Less
// INLINE:   yield
// INLINE: ^bb0([[ARGS]]):
// INLINE:   %cst_2 = arith.constant
// INLINE:   yield
// INLINE-NOT: while_body
// INLINE-NOT: while_cond

// CANON-LABEL:   func @while_main
// CANON-SAME:         ([[VAL_0:%.*]]: tensor<?x256x256xf32>)
// CANON-SAME:         (tensor<i32>, tensor<256x256xf32>, tensor<?x256x256xf32>)
// CANON:           [[VAL_1:%.*]] = arith.constant dense<1.000000e+00> : tensor<256x256xf32>
// CANON:           [[VAL_2:%.*]] = arith.constant dense<0> : tensor<i32>
// CANON:           [[VAL_6:%.*]]:3 = "tfl.while"([[VAL_2]], [[VAL_2]], [[VAL_0]]) <{is_stateless = true}> ({
// CANON:           ^bb0([[VAL_7:%.*]]: tensor<*xi32>, [[VAL_8:%.*]]: tensor<*xi32>, [[VAL_9:%.*]]: tensor<*xf32>):
// CANON:             [[VAL_3:%.*]] = arith.constant dense<10> : tensor<i32>
// CANON:             [[VAL_10:%.*]] = "tf.Less"([[VAL_8]], [[VAL_3]])
// CANON:             "tfl.yield"([[VAL_10]]) : (tensor<*xi1>) -> ()
// CANON:           },  {
// CANON:           ^bb0([[VAL_11:%.*]]: tensor<*xi32>, [[VAL_12:%.*]]: tensor<*xi32>, [[VAL_13:%.*]]: tensor<*xf32>):
// CANON-DAG:         [[VAL_4:%.*]] = arith.constant dense<1> : tensor<i32>
// CANON-DAG:         [[VAL_5:%.*]] = "tf.Const"() <{value = dense<2.560000e+02> : tensor<256x256xf32>}> : () -> tensor<?x?xf32>
// CANON:             [[VAL_14:%.*]] = "tf.AddV2"([[VAL_12]], [[VAL_4]])
// CANON:             [[VAL_15:%.*]] = "tf.AddV2"([[VAL_13]], [[VAL_5]])
// CANON:             [[VAL_16:%.*]] = "tf.AddV2"([[VAL_11]], [[VAL_4]])
// CANON:             "tfl.yield"([[VAL_16]], [[VAL_14]], [[VAL_15]]) : (tensor<*xi32>, tensor<*xi32>, tensor<*xf32>) -> ()
// CANON:           }) : (tensor<i32>, tensor<i32>, tensor<?x256x256xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x256x256xf32>)
// CANON:           return [[VAL_17:%.*]]#1, [[VAL_1]], [[VAL_17]]#2 : tensor<i32>, tensor<256x256xf32>, tensor<?x256x256xf32>
// CANON:         }