1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import RoleFlow.Recipes
import RoleFlow.Search

/-! # The argument-swapping front end

A swappable operation (`Swap`) has two arguments of one sort and a semi-inverse that solves the
first; in the flow signature, `add`. A program that adds `a + b` where `b` is to be consumed is
not typed as written; `swapArgs` rewrites it to `add(b, a)`.

The choice is made by forward inference. Every line contributes its forward form, its solved
form if it has a semi-inverse, and, for a swappable `o(a, b)`, the swapped solved form
`b ← (c, a)`. Starting from the weights and the returned variables, an equation fires when its
reads are known and its results unknown, until nothing fires (`fwdKnown`). The pass swaps
exactly the lines whose swapped form fired. The swap preserves the meaning for a model where
the swappable operations commute (`FrontSem`), and if any swap of the program is typed, the
one chosen here is (`FrontComplete`). -/

namespace RoleFlow

open Search

variable {S : Type}

/-- The swappable operations of a signature: two arguments of one sort, and a semi-inverse that
solves the first. -/
class Swap (Sg : Sig S) where
  swappable : ∀ {as bs : List S}, Sg.Op as bs → Bool
  sorts : ∀ {as bs : List S} (o : Sg.Op as bs), swappable o → ∃ a, as = [a, a]
  k_one : ∀ {as bs : List S} (o : Sg.Op as bs), swappable o → ∃ si, Sg.inv o = some si ∧ si.k = 1

variable {Sg : Sig S} [Swap Sg]

/-- How many leading arguments the semi-inverse of `l` solves, if there is one. -/
def kOf (l : Line Sg) : Option ℕ := (Sg.inv l.op).map (·.k)

/-- The solved form of the line `l`, numbered `i`, at scope `n`. -/
def solvedEqns (i n : ℕ) (l : Line Sg) : List Eqn :=
  match kOf l with
  | some k => [⟨l.args.take k, l.ys n ++ l.args.drop k, i, false⟩]
  | none => []

/-- The swapped solved form: a swappable `o(a, b)` solved for `b`. -/
def swappedEqns (i n : ℕ) (l : Line Sg) : List Eqn :=
  match Swap.swappable l.op, l.args with
  | true, [a, b] => [⟨[b], l.ys n ++ [a], i, true⟩]
  | _, _ => []

/-- The equations of the line `l`: forward, solved, and swapped solved. -/
def lineEqns (i n : ℕ) (l : Line Sg) : List Eqn :=
  ⟨l.ys n, l.args, i, false⟩ :: solvedEqns i n l ++ swappedEqns i n l

/-- The equations of the lines at scope `n`, numbered from `i`. -/
def eqns : ℕ → ℕ → List (Line Sg) → List Eqn
  | _, _, [] => []
  | i, n, l :: ls => lineEqns i n l ++ eqns (i + 1) (n + l.bs.length) ls

/-- Swap the arguments of `l` if it is swappable and `s` holds. -/
def swapIf (s : Bool) (l : Line Sg) : Line Sg :=
  if s && Swap.swappable l.op && l.args.length == 2 then { l with args := l.args.reverse } else l

/-- Swap the lines whose number (from `i`) satisfies `sw`. -/
def swapLines (sw : ℕ → Bool) : ℕ → List (Line Sg) → List (Line Sg)
  | _, [] => []
  | i, l :: ls => swapIf (sw i) l :: swapLines sw (i + 1) ls

/-- The lines whose swapped form fired. -/
def swapSet (P : Prog Sg) : Std.HashSet ℕ :=
  let E := eqns 0 P.ctx₀.length P.lines
  let K₀ := insertAll ∅ (List.range P.wts.length ++ P.ret)
  Std.HashSet.ofList (((fwdKnown E (E.length + 1) K₀ []).2.filter (·.swap)).map (·.line))

/-- The front end: swap the arguments of the lines whose swapped solved form fired. -/
def swapArgs (P : Prog Sg) : Prog Sg :=
  { P with lines := swapLines (swapSet P).contains 0 P.lines }

end RoleFlow