use int.Int
use int.Fact
use int.ComputerDivision
use bintree.Tree
use bintree.Occ as T
use seq.Seq
use seq.FreeMonoid
use seq.Mem as S

use coma.Std

-------------------------------------------------------------------------------

let unTree < 'a > (t: tree 'a)
  (onNode (x: 'a) (l r: tree 'a) {t = Node l x r})
  (onLeaf {t = Empty})
  = any

let get (w: seq int) (i: int) {0 <= i < length w} (k (wi: int) { wi = w[i] })
  = any

-------------------------------------------------------------------------------

let factorial (n: int) {n >= 0} (return (r: int) { r = fact n }) =
  loop {1} {n}
  [ loop (r: int) (k: int) { 0 <= k <= n } { r * (fact k) = fact n }
    = if {k > 0} ((! loop {r * k} {k - 1}))
                 (return {r})
  ]

-------------------------------------------------------------------------------

let product (a b: int) {b >= 0} (return (c: int) { c = a * b })
= loop {a} {b} {0}
  [ loop (p q r: int) { p * q + r = a * b } {q >= 0}
   = if {q > 0} (-> if {mod q 2 = 1}
                  (-> loop {p+p} {div q 2} {r+p})
                  (-> loop {p+p} {div q 2} {r}))
                (-> return {r}) ]

-------------------------------------------------------------------------------

let mergeTree < 'a > (l r: tree 'a)
  (out (o: tree 'a) { forall e. T.mem e o <-> (T.mem e l || T.mem e r) })
= any

let remove_root < 'a > (t: tree 'a)
      (return (o: tree 'a))
= unTree < 'a > {t}
    (fun (x: 'a) (l r: tree 'a) ->
      (! mergeTree < 'a > {l} {r} out)
      [ out (s: tree 'a) { forall e. T.mem e s <-> (T.mem e l || T.mem e r) }
        = return {s} ])
    fail

let remove_root_nobar < 'a > (t: tree 'a)
      (return (o: tree 'a))
= unTree < 'a > {t}
    (fun (x: 'a) (l r: tree 'a) -> mergeTree < 'a > {l} {r} return)
    fail

let client_1 < 'a > (t: tree 'a) {t <> Empty} (return (o: tree 'a) {} )
= remove_root < 'a > {t} return

let client_2 < 'a > (t: tree 'a) (x: 'a) {} (return (o: tree 'a) {} )
= remove_root < 'a > {Node t x t} return

let client_non_provable < 'a > (t: tree 'a) {} (return (o: tree 'a) {} )
= remove_root < 'a > {t} return

-------------------------------------------------------------------------------

-- no_barrier
let find_0 (x: int) (w: seq int) (i: int) (found (ii: int)) (notfound)
= if { i = length w } notfound
     (get {w} {i} (fun (c: int) ->
      if {c=x} (found {i})
               (find_0 {x} {w} {i+1} found notfound)))

let find_1 (x: int) (w: seq int) (i: int)
  (found (ii: int) {i <= ii < length w} {w[ii] = x}) (notfound {not S.mem x w[i..]})
= {0 <= i <= length w}
  (! if {i = length w} notfound
     (get {w} {i} (fun (c: int) ->
      if {c=x} (found {i})
               (find_1 {x} {w} {i+1} found notfound))))

let find_2 (x: int) (w: seq int) (i: int)
  (found (ii: int) {i <= ii < length w} {w[ii] = x}) (notfound {not S.mem x w[i..]})
= {0 <= i <= length w}
  if {i = length w} notfound
     (get {w} {i} (fun (c: int) -> (!
      if {c=x} (found {i})
               (find_2 {x} {w} {i+1} found notfound))))

let find_3 (x: int) (w: seq int) (i: int)
  (found (ii: int) {i <= ii < length w} {w[ii] = x}) (notfound {not S.mem x w[i..]})
= {0 <= i <= length w}
  if {i = length w} notfound
     (get {w} {i} (fun (c: int) ->
      if {c=x} (found {i})
               ((!find_3 {x} {w} {i+1} found notfound))))

-- IMPOSSIBLE to call find_0 under an barrier
-- let fclient_0 (x: int) (w: seq int) (i: int) (found (ii: int) {}) (notfound {})
-- = (! find_0 {x} {w} {i} (fun (ii:int) -> {false} found {ii}) notfound )

let fclient_1 (x: int) (w: seq int) (i: int) {0 <= i <= length w} (found (ii: int) {}) (notfound {})
= find_1 {x} {w} {i} found notfound

let fclient_2 (x: int) (w: seq int) (i: int) {0 <= i <= length w} (found (ii: int) {}) (notfound {})
= find_2 {x} {w} {i} found notfound

let fclient_3 (x: int) (w: seq int) (i: int) {0 <= i <= length w} (found (ii: int) {}) (notfound {})
= find_3 {x} {w} {i} found notfound

-------------------------------------------------------------------------------

let binary_search (s: seq int) (v: int)
  { forall i j. 0 <= i <= j < length s -> s[i] <= s[j] }
    (ret (i: int) { 0 <= i < length s /\ s[i] = v })
    (notfound { forall i. 0 <= i < length s -> s[i] <> v })
= loop {0} {length s}
  [ loop (lo: int) (hi: int)
    { 0 <= lo } { hi <= length s }
    { forall i. 0 <= i < lo -> s[i] <> v }
    { forall i. hi <= i < length s -> s[i] <> v }
    = body {lo + div (hi - lo) 2}
    [ body (mid: int) =
      if {lo >= hi} notfound (-> get {s} {mid} (fun (x: int) ->
      if {v < x} (loop {lo} {mid}) (->
      if {v > x} (loop {mid+1} {hi}) (-> ret {mid}))))
    ]
  ]
