{-# OPTIONS --safe --cubical #-}

module Spartan6.Hierarchy.ResourcePairFlattening where

open import Spartan6.Prelude
open import Spartan6.Hierarchy.Interface
open import Spartan6.Semantics.StateComposition
open import Spartan6.Semantics.StateResource

import Spartan6.Hierarchy.Component as Component
import Spartan6.Hierarchy.FlatMachine as Flat
import Spartan6.Hierarchy.Provenance as Provenance
import Spartan6.Semantics.Machine as Machine

record PairBoundary
  (Input Output : Interface)
  (External LeftObservation RightObservation : Type₀) : Type₀ where
  constructor pairBoundary
  field
    decodePairInput : Environment Input -> External
    encodePairObservation :
      LeftObservation × RightObservation -> Environment Output

open PairBoundary public

private
  variable
    Initial₁ Input₁ Observation₁ State₁ : Type₀
    Initial₂ Input₂ Observation₂ State₂ : Type₀
    External Domain Contract : Type₀
    policy : CrossDomainPolicy Domain Contract

record PairInitialization
  (pair : CoupledResourcePair
    Initial₁ Input₁ Observation₁ State₁
    Initial₂ Input₂ Observation₂ State₂
    External Domain Contract policy) : Type₀ where
  constructor pairInitialization
  field
    pairInitialParameters : Initial₁ × Initial₂
    pairSemanticInitial : State₁ × State₂
    pairInitializationAccepted :
      decodePairInitial pair pairInitialParameters
      ≡ just pairSemanticInitial

open PairInitialization public

pairComponent :
  ∀ {Input Output
    Initial₁ Input₁ Observation₁ State₁
    Initial₂ Input₂ Observation₂ State₂
    External Domain Contract policy}
  -> (identity : Provenance.ComponentIdentity)
  -> (origin : Provenance.SourceOrigin)
  -> (pair : CoupledResourcePair
      Initial₁ Input₁ Observation₁ State₁
      Initial₂ Input₂ Observation₂ State₂
      External Domain Contract policy)
  -> PairBoundary Input Output External Observation₁ Observation₂
  -> PairInitialization pair
  -> Component.Component Input Output (EdgeSet Domain) (State₁ × State₂)
pairComponent identity origin pair boundary initialization =
  Component.leafWithIdentity identity origin
    (Machine.machine
      (pairSemanticInitial initialization)
      (λ input state ->
        encodePairObservation boundary
          (observePair pair (decodePairInput boundary input) state))
      (λ edges input state ->
        stepPair pair edges (decodePairInput boundary input) state))

pairFlatMachine :
  ∀ {Input Output
    Initial₁ Input₁ Observation₁ State₁
    Initial₂ Input₂ Observation₂ State₂
    External Domain Contract policy}
  -> (left-width right-width : ℕ)
  -> (identity : Provenance.ComponentIdentity)
  -> (origin : Provenance.SourceOrigin)
  -> (pair : CoupledResourcePair
      Initial₁ Input₁ Observation₁ State₁
      Initial₂ Input₂ Observation₂ State₂
      External Domain Contract policy)
  -> (boundary : PairBoundary
      Input Output External Observation₁ Observation₂)
  -> (initialization : PairInitialization pair)
  -> CertifiedBitLowering (leftResource pair) left-width
  -> CertifiedBitLowering (rightResource pair) right-width
  -> Flat.FlatMachine Input Output (EdgeSet Domain)
pairFlatMachine
  {Input = Input} {Output = Output}
  {State₁ = State₁} {State₂ = State₂} {Domain = Domain}
  left-width right-width identity origin pair boundary
  initialization left-lowering right-lowering =
  Flat.flatMachine identity
    (Provenance.leafProvenance origin) (left-width + right-width)
    (Machine.machine
      (appendVec
        (encodeState left-lowering (fst initial-state))
        (encodeState right-lowering (snd initial-state)))
      observe-flat
      step-flat)
  where
  initial-state = pairSemanticInitial initialization

  pieces : Vec Bit (left-width + right-width)
    -> Vec Bit left-width × Vec Bit right-width
  pieces = splitVec left-width right-width

  decoded-before : Vec Bit (left-width + right-width) -> State₁ × State₂
  decoded-before state =
    decodeState left-lowering (fst (pieces state))
    , decodeState right-lowering (snd (pieces state))

  observe-flat : Environment Input
    -> Vec Bit (left-width + right-width) -> Environment Output
  observe-flat input state =
    encodePairObservation boundary
      ( observeBits left-lowering left-input (fst (pieces state))
      , observeBits right-lowering right-input (snd (pieces state)))
    where
    external = decodePairInput boundary input
    before = decoded-before state
    left-input = leftInput pair external before
    right-input = rightInput pair external before

  step-flat : EdgeSet Domain -> Environment Input
    -> Vec Bit (left-width + right-width)
    -> Vec Bit (left-width + right-width)
  step-flat edges input state =
    appendVec
      (stepBits left-lowering edges left-input (fst (pieces state)))
      (stepBits right-lowering edges right-input (snd (pieces state)))
    where
    external = decodePairInput boundary input
    before = decoded-before state
    left-input = leftInput pair external before
    right-input = rightInput pair external before

pairFlattens :
  ∀ {Input Output
    Initial₁ Input₁ Observation₁ State₁
    Initial₂ Input₂ Observation₂ State₂
    External Domain Contract policy}
  (left-width right-width : ℕ)
  (identity : Provenance.ComponentIdentity)
  (origin : Provenance.SourceOrigin)
  (pair : CoupledResourcePair
    Initial₁ Input₁ Observation₁ State₁
    Initial₂ Input₂ Observation₂ State₂
    External Domain Contract policy)
  (boundary : PairBoundary Input Output External Observation₁ Observation₂)
  (initialization : PairInitialization pair)
  (left-lowering : CertifiedBitLowering (leftResource pair) left-width)
  (right-lowering : CertifiedBitLowering (rightResource pair) right-width)
  -> Flat.CertifiedMachineFlattening
      (pairComponent identity origin pair boundary initialization)
Flat.flatTarget
  (pairFlattens left-width right-width identity origin pair boundary
    initialization left-lowering right-lowering) =
  pairFlatMachine left-width right-width identity origin pair boundary
    initialization left-lowering right-lowering
Flat.StateRelation
  (pairFlattens left-width right-width identity origin pair boundary
    initialization left-lowering right-lowering)
  semantic-state flat-state =
  (fst (splitVec left-width right-width flat-state)
    ≡ encodeState left-lowering (fst semantic-state))
  ×
  (snd (splitVec left-width right-width flat-state)
    ≡ encodeState right-lowering (snd semantic-state))
Flat.initialRelated
  (pairFlattens left-width right-width identity origin pair boundary
    initialization left-lowering right-lowering) =
  cong fst split-initial , cong snd split-initial
  where
  initial-state = pairSemanticInitial initialization
  split-initial = splitVec-append
    (encodeState left-lowering (fst initial-state))
    (encodeState right-lowering (snd initial-state))
Flat.observationPreserved
  (pairFlattens left-width right-width identity origin pair boundary
    initialization left-lowering right-lowering)
  input semantic-state flat-state related =
  cong (encodePairObservation boundary)
    (cong₂ _,_ left-observation right-observation)
  where
  pieces = splitVec left-width right-width flat-state
  left-bits = fst pieces
  right-bits = snd pieces
  left-related = fst related
  right-related = snd related

  decoded-related :
    (decodeState left-lowering left-bits
    , decodeState right-lowering right-bits)
    ≡ semantic-state
  decoded-related =
    cong₂ _,_
      (cong (decodeState left-lowering) left-related
        ∙ decode-encode left-lowering (fst semantic-state))
      (cong (decodeState right-lowering) right-related
        ∙ decode-encode right-lowering (snd semantic-state))

  external = decodePairInput boundary input
  decoded-state =
    decodeState left-lowering left-bits
    , decodeState right-lowering right-bits

  left-input-related :
    leftInput pair external decoded-state
    ≡ leftInput pair external semantic-state
  left-input-related = cong (leftInput pair external) decoded-related

  right-input-related :
    rightInput pair external decoded-state
    ≡ rightInput pair external semantic-state
  right-input-related = cong (rightInput pair external) decoded-related

  left-observation :
    observeResource (leftResource pair)
      (leftInput pair external semantic-state) (fst semantic-state)
    ≡ observeBits left-lowering
        (leftInput pair external decoded-state) left-bits
  left-observation =
    sym (observe-preserved left-lowering
      (leftInput pair external semantic-state) (fst semantic-state))
    ∙ cong₂ (observeBits left-lowering)
        (sym left-input-related) (sym left-related)

  right-observation :
    observeResource (rightResource pair)
      (rightInput pair external semantic-state) (snd semantic-state)
    ≡ observeBits right-lowering
        (rightInput pair external decoded-state) right-bits
  right-observation =
    sym (observe-preserved right-lowering
      (rightInput pair external semantic-state) (snd semantic-state))
    ∙ cong₂ (observeBits right-lowering)
        (sym right-input-related) (sym right-related)
Flat.transitionPreserved
  (pairFlattens left-width right-width identity origin pair boundary
    initialization left-lowering right-lowering)
  edges input semantic-state flat-state related =
  cong fst split-next ∙ left-next-related
  , cong snd split-next ∙ right-next-related
  where
  pieces = splitVec left-width right-width flat-state
  left-bits = fst pieces
  right-bits = snd pieces
  left-related = fst related
  right-related = snd related

  decoded-state =
    decodeState left-lowering left-bits
    , decodeState right-lowering right-bits

  decoded-related : decoded-state ≡ semantic-state
  decoded-related =
    cong₂ _,_
      (cong (decodeState left-lowering) left-related
        ∙ decode-encode left-lowering (fst semantic-state))
      (cong (decodeState right-lowering) right-related
        ∙ decode-encode right-lowering (snd semantic-state))

  external = decodePairInput boundary input
  left-input-related :
    leftInput pair external decoded-state
    ≡ leftInput pair external semantic-state
  left-input-related = cong (leftInput pair external) decoded-related
  right-input-related :
    rightInput pair external decoded-state
    ≡ rightInput pair external semantic-state
  right-input-related = cong (rightInput pair external) decoded-related

  next-left-bits =
    stepBits left-lowering edges
      (leftInput pair external decoded-state) left-bits
  next-right-bits =
    stepBits right-lowering edges
      (rightInput pair external decoded-state) right-bits

  split-next = splitVec-append next-left-bits next-right-bits

  left-next-related :
    next-left-bits
    ≡ encodeState left-lowering
        (fst (stepPair pair edges external semantic-state))
  left-next-related =
    cong₂ (stepBits left-lowering edges)
      left-input-related left-related
    ∙ step-lowering-preserved left-lowering edges
        (leftInput pair external semantic-state) (fst semantic-state)

  right-next-related :
    next-right-bits
    ≡ encodeState right-lowering
        (snd (stepPair pair edges external semantic-state))
  right-next-related =
    cong₂ (stepBits right-lowering edges)
      right-input-related right-related
    ∙ step-lowering-preserved right-lowering edges
        (rightInput pair external semantic-state) (snd semantic-state)
Flat.provenanceRetained
  (pairFlattens left-width right-width identity origin pair boundary
    initialization left-lowering right-lowering) = refl