import Set
import Option

fun operator ∘ <T,U,V>(g:fn U->V, f:fn T->U) {
  fun x:T { g(f(x)) }
}

fun update<T, U>(f:fn T->U, x:T, v:U) {
  fun y:T { if y = x then v else f(y) }
}

theorem update_eq :
  < T, U > all f: fn (T) -> U, x:T, v:U.
  update(f, x, v)(x) = v
proof
  arbitrary T:type, U:type
  arbitrary f: fn (T) -> U, x:T, v:U
  expand update.
end

theorem update_not_eq :
  < T, U > all f: fn (T) -> U, x:T, v:U, y:T.
  if not (x = y)
  then update(f, x, v)(y) = f(y)
proof
  arbitrary T:type, U:type
  arbitrary f: fn (T) -> U, x:T, v:U, y:T
  suppose not_eq
  switch y = x for update {
    case true suppose yx_true {
      have yx: y = x by replace yx_true.
      have xy: x = y by symmetric yx
      conclude false by apply not_eq to xy
    }
    case false {
      .
    }
  }
end

theorem update_shadow : all T:type, U:type. all f:fn(T)->U, x:T, v:U, w:U.
  update(update(f, x, v), x, w) = update(f, x, w)
proof
  arbitrary T:type, U:type
  arbitrary f:fn(T)->U, x:T, v:U, w:U
  extensionality
  arbitrary y:T
  switch y = x for update {
    case true { . }
    case false { . }
  }
end

theorem update_permute :
  < T, U > all f:fn(T)->U, x:T, v:U, w:U, y:T.
  if not (x = y)
  then update(update(f, x, v), y, w) = update(update(f, y, w), x, v)
proof
  arbitrary T:type, U:type
  arbitrary f:fn(T)->U, x:T, v:U, w:U, y:T
  suppose x_neq_y
  extensionality
  arbitrary z:T
  switch z = y for update {
    case true suppose zy_true {
      have zy: z = y by replace zy_true.
      switch z = x {
        case true suppose zx_true {
	        have zx: z = x by replace zx_true.
	        have xy: x = y by transitive (symmetric zx) zy
	        conclude false by apply x_neq_y to xy
	      }
	      case false {
	        .
	      }
      }
    }
    case false {
      .
    }
  }
end

/* Partial Maps */

define empty_map = generic T, U { λk:T { @none<U> } }

/*
define domain : fn  (fn(T)->Option) -> Set
              = generic T, U { λf: fn(T)->Option {char_fun(λx:T{
                  switch f(x) {
	            case none { false }
		    case just(y) { true }
		  }})}}
*/
fun domain<T,U>(f : fn(T)->Option<U>) {
  set_of_pred(fun x:T {
      switch f(x) {
        case none { false }
        case just(y) { true }
      }})
}


define restrict = generic T, U { λf:fn(T)->Option<U>, P:Set<T> { λx:T {
                      if x ∈ P then f(x) else @none<U> }}}

define combine = generic T, U { fun f:fn(T)->Option<U>, g:fn(T)->Option<U> { fun x:T {
                     if f(x) = none then g(x) else f(x) }}}

theorem combine_left: all T:type, U:type, f:fn(T)->Option<U>, g:fn(T)->Option<U>, x:T.
  if g(x) = none
  then combine(f, g)(x) = f(x)
proof
  arbitrary T:type, U:type
  arbitrary f:fn(T)->Option<U>, g:fn(T)->Option<U>, x:T
  assume: g(x) = none
  switch f(x) {
    case none assume: f(x) = none {
      show combine(f, g)(x) = f(x)
      expand combine
      replace (recall f(x) = none) | (recall g(x) = none).
    }
    case just(v) assume: f(x) = just(v) {
      show combine(f, g)(x) = f(x)
      expand combine
      replace recall f(x) = just(v).
    }
  }
end

theorem combine_right: all T:type, U:type, f:fn(T)->Option<U>, g:fn(T)->Option<U>, x:T.
  if f(x) = none
  then combine(f, g)(x) = g(x)
proof
  arbitrary T:type, U:type
  arbitrary f:fn(T)->Option<U>, g:fn(T)->Option<U>, x:T
  assume: f(x) = none
  show combine(f, g)(x) = g(x)
  expand combine
  replace recall f(x) = none.
end

theorem restrict_domain:
  < T, U > all f:fn(T)->Option<U>, P:Set<T>.
  domain(restrict(f, P)) ⊆ P
proof
  arbitrary T:type, U:type
  arbitrary f:fn(T)->Option<U>, P:Set<T>
  expand restrict | domain | operator ⊆
  arbitrary y:T
  switch y ∈ P {
    case true { . }
    case false suppose Px_false {
      assume prem
      conclude false by replace Px_false in apply member_set_of_pred to prem
    }
  }
end

define flip : fn<T,U,V> (fn T,U->V) ->(fn U,T->V)
  = generic T,U,V { λf{ λx,y{ f(y,x) }}}

theorem flip_flip:
  all T:type. all f:fn T,T->T. flip(flip(f)) = f
proof
  arbitrary T:type
  arbitrary f:fn T,T->T
  extensionality
  arbitrary x:T, y:T
  expand flip.
end