open import LFT
module ST (Iota : Set) where

open import Relation.Binary.PropositionalEquality
open import Function
open import Data.Unit using (⊤)
open import Data.Nat
open import Reflection
open import Coinduction

data Free (S : Set) (T : S -> Set) (A : Set) : Set where
  ret : A -> Free S T A
  _<<_ : (s : S) -> (T s -> Free S T A) -> Free S T A

data STS (s : Set) : Set where
  new : (x : Iota) -> STS s
  read : (ref : s) -> STS s
  write : (ref : s) -> (x : Iota) -> STS s

STT : ∀ {s} -> STS s -> Set
STT {s} (new _) = s
STT (read _) = Iota
STT (write _ _) = ⊤

ST : (s a : Set) -> Set
ST s a = Free (STS s) STT a

data SafeS (n : ℕ) : STS ℕ -> Set where
  new : (x : Iota) -> SafeS n (new x) 
  read : {ref : ℕ} -> ref < n -> SafeS n (read ref)
  write : {ref : ℕ} -> ref < n -> (x : Iota) -> SafeS n (write ref x)

SafeT : (n : ℕ) -> (s : STS ℕ) -> STT s -> Set
SafeT n (new x) m = n ≡ m
SafeT n (read ref) i = ⊤
SafeT n (write ref x) _ = ⊤

next : ℕ -> STS ℕ -> ℕ
next n (new x) = suc n
next n (read ref) = n
next n (write ref x) = n

data IFree {S T A} {I : Set} (IS : I -> S -> Set)(IT : I -> (s : S) -> T s -> Set)(r : I -> S -> I) (i : I) : Free S T A -> Set where
  ret : (a : A) -> IFree IS IT r i (ret a)
  _<<_ : ∀ {s t} -> IS i s -> (∀ x -> IT i s x -> IFree IS IT r (r i s) (t x)) -> IFree IS IT r i (s << t)

Safe : ∀ {a} -> ℕ -> ST ℕ a -> Set
Safe = IFree SafeS SafeT next

[Iota] : [Set] Iota Iota
[Iota] = _≡_

[⊤] : [Set] ⊤ ⊤
[⊤] = \ _ _ -> ⊤

data [Free] {S1 S2 : Set}([S] : [Set] S1 S2) {T1 : S1 -> Set}{T2 : S2 -> Set}([T] : ([S] [->] [Set]) T1 T2) {A1 A2 : Set}([A] : [Set] A1 A2) 
     : [Set] (Free S1 T1 A1) (Free S2 T2 A2) where
  [ret] : ([A] [->] [Free] [S] [T] [A]) ret ret
  _[<<]_ : ([Π] [S] \ [s] ->  ([T] [s] [->] [Free] [S] [T] [A]) [->] [Free] [S] [T] [A]) _<<_ _<<_

data [STS] {s1 s2 : Set}([s] : [Set] s1 s2) : [Set] (STS s1) (STS s2) where
  [new] : ([Iota] [->] [STS] [s]) new new
  [read] : ([s] [->] [STS] [s]) read read
  [write] : ([s] [->] [Iota] [->] [STS] [s]) write write

[STT] : ∀ {s1 s2} {[s] : [Set] s1 s2} -> ([STS] [s] [->] [Set]) STT STT
[STT] {[s] = [s]} ([new] aR) = [s]
[STT] ([read] aR) = [Iota]
[STT] ([write] aR aR') = [⊤]

[ST] : ([Set] [->] [Set] [->] [Set]) ST ST
[ST] [s] [a] = [Free] ([STS] [s]) [STT] [a]

postulate
  free-prog : ∀ {a} -> (prog : ∀ s -> ST s a) -> ([Π] [Set] \ [s] -> [ST] [s] _≡_) prog prog

C : {a : Set} -> (n : ℕ) -> ST ℕ a -> Set
C {a} n m = [ST] (λ i m → m < n) _≡_ m m

data Stream (A : Set) (n : ℕ) (m : ST ℕ A) : Set where
  _∷_ : (x : C {A} n m) (xs : ∞ (Stream A (suc n) m)) -> Stream A n m

≤-refl : ∀ {m n : ℕ} -> n ≡ m -> m < suc n
≤-refl {zero} eq = s≤s z≤n
≤-refl {suc n} refl = s≤s (≤-refl refl)

≤-suc : ∀ {m n} -> m ≤ n -> m ≤ suc n
≤-suc z≤n = z≤n
≤-suc (s≤s m≤n) = s≤s (≤-suc m≤n)

shift[read] : ∀ {a n a1 a2 x} -> Stream a n (read a1 << a2) -> Stream a n (a2 x)
shift[read] (([read] _ [<<] x) ∷ xs) = x refl ∷ (♯ shift[read] (♭ xs))

shift[write] : ∀ {a n a1 x a2} -> Stream a n (write a1 x << a2) -> Stream a n (a2 _)
shift[write] (([write] _ _ [<<] x) ∷ xs) = x _ ∷ (♯ shift[write] (♭ xs))

shift[new] : ∀ {a n x a1 a2} -> x < n -> Stream a n (new a1 << a2) -> Stream a n (a2 x)
shift[new] lt (([new] aR [<<] aR') ∷ xs) = aR' lt ∷ ♯ shift[new] (≤-suc lt) (♭ xs)

helper : ∀ {a} n {m2} -> C {a} n m2 -> Stream a (suc n) m2 -> Safe n m2
helper n ([ret] aR) _ = ret _
helper n ([new] _ [<<] _) (([new] _ [<<] aR) ∷ xs) = new _ << λ x n≡x → helper (suc n) (aR (≤-refl n≡x)) (shift[new] (≤-suc (≤-refl n≡x)) (♭ xs))
helper n ([read] aR [<<] aR') xs = read aR << λ x _ → helper n (aR' refl) (shift[read] xs)
helper n ([write] aR aR' [<<] aR0) xs = write aR _ << λ x _ → helper n (aR0 _) (shift[write] xs)

mkStream : ∀ {a} n -> (m : ∀ s -> ST s a) -> Stream a n (m ℕ)
mkStream n m = free-prog m _ ∷ (♯ (mkStream (suc n) m))

theorem : ∀ {a} -> (m : ∀ s -> ST s a) -> Safe 0 (m ℕ)
theorem m with mkStream 0 m 
theorem m | x ∷ xs = helper 0 x (♭ xs)