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

module Spartan6.Hierarchy.FeedbackFlattening where

open import Spartan6.Prelude
open import Spartan6.Hierarchy.Interface

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

-- Feedback becomes executable flat state only through this explicit layout
-- contract.  In particular, constructing a guarded hierarchy alone does not
-- claim that its retained feedback environment has a bit-level allocation.

record FeedbackAllocation (Feedback : Interface) : Type₀ where
  constructor feedbackAllocation
  field
    feedbackStateCount : ℕ
    encodeFeedback :
      Environment Feedback -> Vec Bit feedbackStateCount
    decodeFeedback :
      Vec Bit feedbackStateCount -> Environment Feedback
    decode-encode-feedback : ∀ feedback
      -> decodeFeedback (encodeFeedback feedback) ≡ feedback

open FeedbackAllocation public

interfaceFeedbackAllocation : (Feedback : Interface)
  -> FeedbackAllocation Feedback
feedbackStateCount (interfaceFeedbackAllocation Feedback) =
  interfaceWidth Feedback
encodeFeedback (interfaceFeedbackAllocation Feedback) = flattenEnvironment
decodeFeedback (interfaceFeedbackAllocation Feedback) = unflattenEnvironment
decode-encode-feedback (interfaceFeedbackAllocation Feedback) =
  unflatten-flatten

feedbackFlatMachine :
  ∀ {ExternalInput ExternalOutput Feedback Event State}
  (identity : Provenance.ComponentIdentity)
  (guard : Feedback.FeedbackGuard Event)
  (initial-feedback : Environment Feedback)
  (allocation : FeedbackAllocation Feedback)
  {original : Component.Component
    (ExternalInput ∥ᵢ Feedback)
    (ExternalOutput ∥ᵢ Feedback)
    Event State}
  -> Flat.CertifiedMachineFlattening original
  -> Flat.FlatMachine ExternalInput ExternalOutput Event
feedbackFlatMachine
  {ExternalInput = ExternalInput} {ExternalOutput = ExternalOutput}
  {Feedback = Feedback} {Event = Event} {State = State}
  identity guard initial-feedback allocation {original} certificate =
  Flat.flatMachine identity
    (Provenance.feedbackProvenance
      (Provenance.identityName identity)
      (Component.componentProvenance original))
    (old-count + feedback-count)
    (Machine.machine
      (appendVec
        (Flat.flatInitial original-target)
        (encodeFeedback allocation initial-feedback))
      observe-flat
      step-flat)
  where
  original-target = Flat.flatTarget certificate
  old-count = Flat.flatStateCount original-target
  feedback-count = feedbackStateCount allocation

  pieces : Vec Bit (old-count + feedback-count)
    -> Vec Bit old-count × Vec Bit feedback-count
  pieces = splitVec old-count feedback-count

  observe-flat : Environment ExternalInput
    -> Vec Bit (old-count + feedback-count)
    -> Environment ExternalOutput
  observe-flat input state =
    fst
      (Flat.flatObserve original-target
        (input , decodeFeedback allocation feedback-bits)
        original-bits)
    where
    original-bits = fst (pieces state)
    feedback-bits = snd (pieces state)

  step-flat : Event -> Environment ExternalInput
    -> Vec Bit (old-count + feedback-count)
    -> Vec Bit (old-count + feedback-count)
  step-flat event input state =
    appendVec
      (Flat.flatStep original-target event
        (input , decoded-feedback) original-bits)
      (if Feedback.advancesFeedback guard event
       then encodeFeedback allocation (snd observed)
       else feedback-bits)
    where
    original-bits = fst (pieces state)
    feedback-bits = snd (pieces state)
    decoded-feedback = decodeFeedback allocation feedback-bits
    observed = Flat.flatObserve original-target
      (input , decoded-feedback) original-bits

feedbackFlattens :
  ∀ {ExternalInput ExternalOutput Feedback Event State}
  (identity : Provenance.ComponentIdentity)
  (guard : Feedback.FeedbackGuard Event)
  (initial-feedback : Environment Feedback)
  (allocation : FeedbackAllocation Feedback)
  {original : Component.Component
    (ExternalInput ∥ᵢ Feedback)
    (ExternalOutput ∥ᵢ Feedback)
    Event State}
  (certificate : Flat.CertifiedMachineFlattening original)
  -> Flat.CertifiedMachineFlattening
      (Feedback.guardedFeedbackWithIdentity
        identity guard initial-feedback original)
Flat.flatTarget
  (feedbackFlattens identity guard initial-feedback allocation certificate) =
  feedbackFlatMachine
    identity guard initial-feedback allocation certificate
Flat.StateRelation
  (feedbackFlattens identity guard initial-feedback allocation certificate)
  semantic-state flat-state =
  Flat.StateRelation certificate
    (fst semantic-state)
    (fst (splitVec old-count feedback-count flat-state))
  ×
  (snd (splitVec old-count feedback-count flat-state)
    ≡ encodeFeedback allocation (snd semantic-state))
  where
  old-count = Flat.flatStateCount (Flat.flatTarget certificate)
  feedback-count = feedbackStateCount allocation
Flat.initialRelated
  (feedbackFlattens identity guard initial-feedback allocation
    {original = original} certificate) =
  subst
    (Flat.StateRelation certificate
      (Component.componentInitial original))
    (sym (cong fst split-initial))
    (Flat.initialRelated certificate)
  , cong snd split-initial
  where
  original-target = Flat.flatTarget certificate
  split-initial = splitVec-append
    (Flat.flatInitial original-target)
    (encodeFeedback allocation initial-feedback)
Flat.observationPreserved
  (feedbackFlattens identity guard initial-feedback allocation
    {original = original} certificate)
  input semantic-state flat-state related =
  cong fst
    (Flat.observationPreserved certificate
      (input , semantic-feedback)
      semantic-original original-bits original-related
    ∙ cong
        (λ feedback ->
          Flat.flatObserve original-target
            (input , feedback) original-bits)
        (sym decoded-feedback-related))
  where
  original-target = Flat.flatTarget certificate
  old-count = Flat.flatStateCount original-target
  feedback-count = feedbackStateCount allocation
  pieces = splitVec old-count feedback-count flat-state
  original-bits = fst pieces
  feedback-bits = snd pieces
  semantic-original = fst semantic-state
  semantic-feedback = snd semantic-state
  original-related = fst related
  feedback-related = snd related
  decoded-feedback-related :
    decodeFeedback allocation feedback-bits ≡ semantic-feedback
  decoded-feedback-related =
    cong (decodeFeedback allocation) feedback-related
    ∙ decode-encode-feedback allocation semantic-feedback
Flat.transitionPreserved
  (feedbackFlattens identity guard initial-feedback allocation
    {original = original} certificate)
  event input semantic-state flat-state related =
  subst
    (Flat.StateRelation certificate semantic-original-next)
    (sym (cong fst split-next))
    original-next-related
  , cong snd split-next ∙ feedback-next-related
  where
  original-target = Flat.flatTarget certificate
  old-count = Flat.flatStateCount original-target
  feedback-count = feedbackStateCount allocation
  pieces = splitVec old-count feedback-count flat-state
  original-bits = fst pieces
  feedback-bits = snd pieces
  semantic-original = fst semantic-state
  semantic-feedback = snd semantic-state
  original-related = fst related
  feedback-related = snd related

  decoded-feedback = decodeFeedback allocation feedback-bits
  decoded-feedback-related : decoded-feedback ≡ semantic-feedback
  decoded-feedback-related =
    cong (decodeFeedback allocation) feedback-related
    ∙ decode-encode-feedback allocation semantic-feedback

  semantic-observation =
    Component.componentObserve original
      (input , semantic-feedback) semantic-original
  flat-observation =
    Flat.flatObserve original-target
      (input , decoded-feedback) original-bits

  observation-related : semantic-observation ≡ flat-observation
  observation-related =
    Flat.observationPreserved certificate
      (input , semantic-feedback)
      semantic-original original-bits original-related
    ∙ cong
        (λ feedback ->
          Flat.flatObserve original-target
            (input , feedback) original-bits)
        (sym decoded-feedback-related)

  semantic-original-next =
    Component.componentStep original event
      (input , semantic-feedback) semantic-original
  flat-original-next =
    Flat.flatStep original-target event
      (input , decoded-feedback) original-bits

  original-step-input-related :
    Flat.flatStep original-target event
      (input , semantic-feedback) original-bits
    ≡ flat-original-next
  original-step-input-related =
    cong
      (λ feedback ->
        Flat.flatStep original-target event
          (input , feedback) original-bits)
      (sym decoded-feedback-related)

  original-next-related :
    Flat.StateRelation certificate
      semantic-original-next flat-original-next
  original-next-related =
    subst
      (Flat.StateRelation certificate semantic-original-next)
      original-step-input-related
      (Flat.transitionPreserved certificate
        event (input , semantic-feedback)
        semantic-original original-bits original-related)

  semantic-feedback-next =
    if Feedback.advancesFeedback guard event
    then snd semantic-observation
    else semantic-feedback
  flat-feedback-next =
    if Feedback.advancesFeedback guard event
    then encodeFeedback allocation (snd flat-observation)
    else feedback-bits

  feedback-next-related :
    flat-feedback-next
    ≡ encodeFeedback allocation semantic-feedback-next
  feedback-next-related with Feedback.advancesFeedback guard event
  ... | false = feedback-related
  ... | true =
    cong (λ output -> encodeFeedback allocation (snd output))
      (sym observation-related)

  split-next = splitVec-append
    flat-original-next flat-feedback-next
Flat.provenanceRetained
  (feedbackFlattens identity guard initial-feedback allocation certificate) =
  refl