Graph.lean

import GraphOptVerify.Core.Relation

namespace GraphOptVerify.ONNX

universe u

/-- The node representation is independent of the scalar interpretation.
Port-layout lists describe actual ports, not zero-length variadic sentinels. -/
structure Node where
  schema : SchemaId
  attributes : AttributeValues schema
  inputNames : List (Option String)
  outputNames : List (Option String)
  inputLayout : List PortLayout
  outputLayout : List PortLayout
  captureNames : List String := []
  occurrence : List Nat := []
  deriving Repr

def Node.liveOutputNames (node : Node) : List String :=
  node.outputNames.filterMap id

def Node.WellFormed (node : Node) : Prop :=
  node.inputLayout.length = node.inputNames.length ∧
    node.outputLayout.length = node.outputNames.length ∧
    node.liveOutputNames.Nodup ∧
    node.captureNames.Nodup ∧
    (∀ name ∈ node.inputNames.filterMap id, name ≠ "") ∧
    (∀ name ∈ node.liveOutputNames, name ≠ "") ∧
    (∀ name ∈ node.captureNames, name ≠ "")

structure Graph (model : ValueModel.{u}) where
  shapeAssignment : ShapeAssignment := []
  inputs : List (String × ValueType)
  initializers : List (String × model.Carrier × ValueType)
  nodes : List Node
  outputs : List (String × ValueType)

def resolveNames {model : ValueModel.{u}} (environment : Environment model) :
    List String → Option (List model.Carrier)
  | [] => some []
  | name :: names => do
      let value ← environment.lookup name
      let remaining ← resolveNames environment names
      pure (value :: remaining)

/-- None in the port list is preserved; an unresolved present name fails. -/
def resolveInputs {model : ValueModel.{u}} (environment : Environment model) :
    List (Option String) → Option (InputValues model)
  | [] => some []
  | none :: names => do
      let remaining ← resolveInputs environment names
      pure (none :: remaining)
  | some name :: names => do
      let value ← environment.lookup name
      let remaining ← resolveInputs environment names
      pure (some value :: remaining)

def namesFresh {model : ValueModel.{u}}
    (environment : Environment model) (names : List String) : Prop :=
  ∀ name ∈ names, environment.lookup name = none

/-- NodeEval separately checks the masks, so zip cannot hide a missing port. -/
def bindOutputs {model : ValueModel.{u}} (environment : Environment model)
    (names : List (Option String)) (outputs : OutputValues model) : Environment model :=
  environment ++ (names.zip outputs).filterMap (fun pair =>
    match pair with
    | (some name, some value) => some (name, value)
    | _ => none)

def EnvironmentHasTypes {model : ValueModel.{u}}
    (assignment : ShapeAssignment) (environment : Environment model)
    (declared : List (String × ValueType)) : Prop :=
  ∀ declaration ∈ declared,
    ∃ value, environment.lookup declaration.1 = some value ∧
      model.hasType assignment value declaration.2

def initializerEnvironment {model : ValueModel.{u}} (graph : Graph model) : Environment model :=
  graph.initializers.map (fun initializer => (initializer.1, initializer.2.1))

theorem EnvironmentHasTypes.single {model : ValueModel.{u}}
    {assignment : ShapeAssignment} (environment : Environment model)
    (name : String) (type : ValueType) (value : model.Carrier)
    (resolved : environment.lookup name = some value)
    (typed : model.hasType assignment value type) :
    EnvironmentHasTypes assignment environment [(name, type)] := by
  intro declaration member
  have same : declaration = (name, type) := List.mem_singleton.mp member
  subst declaration
  exact ⟨value, resolved, typed⟩

def initializerTypes {model : ValueModel.{u}} (graph : Graph model) : List (String × ValueType) :=
  graph.initializers.map (fun initializer => (initializer.1, initializer.2.2))

def Node.call {model : ValueModel.{u}} (node : Node) (context : EvalContext model)
    (assignment : ShapeAssignment) (inputs : InputValues model)
    (capturedValues : List model.Carrier) : OperatorCall model :=
  { attributes := node.attributes
    inputs := inputs
    inputLayout := node.inputLayout
    outputMask := node.outputNames.map Option.isSome
    outputLayout := node.outputLayout
    captures := node.captureNames.zip capturedValues
    occurrence := node.occurrence
    shapeAssignment := assignment
    randomTrace := context.randomTrace
    externalProfile := context.externalProfile }

inductive NodeEval {model : ValueModel.{u}} (context : EvalContext model)
    (assignment : ShapeAssignment) : Node → Environment model → Environment model → Prop
  | step
      (node : Node)
      (environment : Environment model)
      (inputs : InputValues model)
      (capturedValues : List model.Carrier)
      (outputs : OutputValues model)
      (wellFormed : node.WellFormed)
      (inputsResolved : resolveInputs environment node.inputNames = some inputs)
      (capturesResolved : resolveNames environment node.captureNames = some capturedValues)
      (outputsFresh : namesFresh environment node.liveOutputNames)
      (outputPorts : outputMask outputs = node.outputNames.map Option.isSome)
      (relation : OpRel node.schema context (node.call context assignment inputs capturedValues) outputs) :
      NodeEval context assignment node environment (bindOutputs environment node.outputNames outputs)

inductive NodesEval {model : ValueModel.{u}} (context : EvalContext model)
    (assignment : ShapeAssignment) :
    List Node → Environment model → Environment model → Prop
  | nil (environment : Environment model) : NodesEval context assignment [] environment environment
  | cons
      (node : Node)
      (nodes : List Node)
      (before middle after : Environment model)
      (head : NodeEval context assignment node before middle)
      (tail : NodesEval context assignment nodes middle after) :
      NodesEval context assignment (node :: nodes) before after

def runtimeBoundaryMatches {model : ValueModel.{u}}
    (graph : Graph model) (runtimeInputs : Environment model) : Prop :=
  runtimeInputs.map Prod.fst = graph.inputs.map Prod.fst

def initializerNamesFresh {model : ValueModel.{u}}
    (graph : Graph model) (runtimeInputs : Environment model) : Prop :=
  namesFresh runtimeInputs (graph.initializers.map (fun initializer => initializer.1))

inductive GraphEval {model : ValueModel.{u}} (context : EvalContext model) :
    Graph model → Environment model → List model.Carrier → Prop
  | run
      (graph : Graph model)
      (runtimeInputs finalEnvironment : Environment model)
      (outputs : List model.Carrier)
      (runtimeBoundary : runtimeBoundaryMatches graph runtimeInputs)
      (inputNamesUnique : (graph.inputs.map Prod.fst).Nodup)
      (initializerNamesUnique : (graph.initializers.map (fun initializer => initializer.1)).Nodup)
      (outputNamesUnique : (graph.outputs.map Prod.fst).Nodup)
      (runtimeInputsTyped : EnvironmentHasTypes graph.shapeAssignment runtimeInputs graph.inputs)
      (initializersFresh : initializerNamesFresh graph runtimeInputs)
      (initializersTyped :
        EnvironmentHasTypes graph.shapeAssignment (initializerEnvironment graph) (initializerTypes graph))
      (nodes : NodesEval context graph.shapeAssignment graph.nodes
        (runtimeInputs ++ initializerEnvironment graph) finalEnvironment)
      (outputsResolved : resolveNames finalEnvironment (graph.outputs.map Prod.fst) = some outputs)
      (outputsTyped : EnvironmentHasTypes graph.shapeAssignment finalEnvironment graph.outputs) :
      GraphEval context graph runtimeInputs outputs

end GraphOptVerify.ONNX