-----Description-----  
The task is to search for a specific integer in a 2D array where the rows and columns are sorted in non-decreasing order. The goal is to locate the key and return its position as row and column indices, or return (-1, -1) if the algorithm fails to find the key.

-----Input-----
The input consists of:
• a: A non-empty 2D array of integers (Array (Array Int)) where:
  - The outer array contains at least one row
  - Each row has the same non-zero length
• key: An integer value (Int) to search for in the array.

-----Output-----  
The output is a pair of integers (Int × Int):  
• If the key is found, the first element represents the row index and the second element represents the column index such that get2d a row col = key.  
• If the key is not found, the function returns (-1, -1).

-----Note-----  
It is assumed that the input 2D array is sorted by rows and columns.

Difficulty: basic

Lean 4 solution:
-- !benchmark @start import type=solution

-- !benchmark @end import

-- !benchmark @start solution_aux
@[reducible, simp]
def get2d (a : Array (Array Int)) (i j : Int) : Int :=
  (a[Int.toNat i]!)[Int.toNat j]!
-- !benchmark @end solution_aux

-- !benchmark @start precond_aux
-- !benchmark @end precond_aux

@[reducible, simp]
def SlopeSearch_precond (a : Array (Array Int)) (key : Int) : Prop :=
  -- !benchmark @start precond
  a.size > 0 ∧
  (a[0]!).size > 0 ∧  -- non-empty inner arrays
  List.Pairwise (·.size = ·.size) a.toList ∧
  a.all (fun x => List.Pairwise (· ≤ ·) x.toList) ∧
  (List.range (a[0]!.size)).all (fun i =>
    List.Pairwise (· ≤ ·) (a.map (fun x => x[i]!)).toList
  )
  -- !benchmark @end precond


-- !benchmark @start code_aux

-- !benchmark @end code_aux


def SlopeSearch (a : Array (Array Int)) (key : Int) (h_precond : SlopeSearch_precond (a) (key)) : (Int × Int) :=
  -- !benchmark @start code
  let rows := a.size
  let cols := if rows > 0 then (a[0]!).size else 0

  let rec aux (m n : Int) (fuel : Nat) : (Int × Int) :=
    if fuel = 0 then (-1, -1)
    else if m ≥ Int.ofNat rows || n < 0 then (-1, -1)
    else if get2d a m n = key then (m, n)
    else if get2d a m n < key then
      aux (m + 1) n (fuel - 1)
    else
      aux m (n - 1) (fuel - 1)

  aux 0 (Int.ofNat (cols - 1)) (rows + cols)
  -- !benchmark @end code


-- !benchmark @start postcond_aux

-- !benchmark @end postcond_aux


@[reducible, simp]
def SlopeSearch_postcond (a : Array (Array Int)) (key : Int) (result: (Int × Int)) (h_precond : SlopeSearch_precond (a) (key)) :=
  -- !benchmark @start postcond
  let (m, n) := result;
  (m ≥ 0 ∧ m < a.size ∧ n ≥ 0 ∧ n < (a[0]!).size ∧ get2d a m n = key) ∨
  (m = -1 ∧ n = -1 ∧ a.all (fun x => x.all (fun e => e ≠ key)))
  -- !benchmark @end postcond


-- !benchmark @start proof_aux

-- !benchmark @end proof_aux


theorem SlopeSearch_spec_satisfied (a: Array (Array Int)) (key: Int) (h_precond : SlopeSearch_precond (a) (key)) :
    SlopeSearch_postcond (a) (key) (SlopeSearch (a) (key) h_precond) h_precond := by
  -- !benchmark @start proof
  sorry
  -- !benchmark @end proof


Tests:
input: a: #[#[1, 2, 3], #[4, 5, 6], #[7, 8, 9]]
key: 5
a: #[#[1, 2, 3], #[4, 5, 6], #[7, 8, 9]]
key: 3
a: #[#[1, 2, 3], #[4, 5, 6], #[7, 8, 9]]
key: 10
a: #[#[1, 2, 3, 4]]
key: 4
a: #[#[1], #[2], #[3], #[4]]
key: 3
expected: (1, 1)
(0, 2)
(-1, -1)
(0, 3)
(2, 0)
unexpected: (1, 2)
(0, 1)
(0, 1)
(1, 2)
(1, 1)
(2, 2)
(0, 2)
(1, 3)
(0, 4)
(1, 0)
(2, 1)

Inputs rejected by the precondition:
a: #[#[1, 3, 2], #[0, 6, 5], #[7, 8, 9]]
key: 2
a: #[]
key: 5
a: #[#[]]
key: 5
a: #[#[], #[]]
key: 5