Map fusion

Definition 1 (Output specification).

For finite lists $x=[x_1,\ldots,x_n]\in\operatorname{List}(\alpha)$ and $y\in\operatorname{List}(\gamma)$, the specification predicate $\mathcal{S}_{f,g}(x,y)$ means that $y$ preserves the length and order of $x$ and replaces each element $x_i$ with $g(f(x_i))$:

In Lean, Specification f g xs ys expresses this relation inductively: nil relates the two empty lists, and cons relates their corresponding heads and tails. Here xs and ys denote $x$ and $y$.

inductive Specification (f : α → β) (g : β → γ) : List α → List γ → Prop
  | nil : Specification f g [] []
  | cons (x : α) (xs : List α) (ys : List γ) :
      Specification f g xs ys → Specification f g (x :: xs) (g (f x) :: ys)

Theorem 2 (Correctness and equivalence).

For every $x\in\operatorname{List}(\alpha)$, both implementations satisfy the specification and produce the same output:

Proof.

Satisfaction. Induction on $x$ proves versionA_correct and versionB_correct. The empty case uses Specification.nil; the nonempty case uses Specification.cons and the induction hypothesis.

theorem versionA_correct (f : α → β) (g : β → γ) (xs : List α) :
    Specification f g xs (versionA f g xs) := by
  induction xs with
  | nil => exact .nil
  | cons x xs ih =>
    simpa only [versionA, List.map_cons] using
      Specification.cons x xs (versionA f g xs) ih

theorem versionB_correct (f : α → β) (g : β → γ) (xs : List α) :
    Specification f g xs (versionB f g xs) := by
  induction xs with
  | nil => exact .nil
  | cons x xs ih =>
    simpa only [versionB, List.map_cons] using
      Specification.cons x xs (versionB f g xs) ih

Uniqueness. The specification determines one output:

theorem Specification.unique {f : α → β} {g : β → γ} {xs : List α}
    {ys zs : List γ} (hy : Specification f g xs ys) (hz : Specification f g xs zs) :
    ys = zs := by
  induction hy generalizing zs with
  | nil => cases hz; rfl
  | cons x xs ys h ih =>
    cases hz with
    | cons _ _ zs hz => exact congrArg (List.cons (g (f x))) (ih hz)

Equivalence. Apply uniqueness to the two satisfaction proofs:

theorem versionA_eq_versionB (f : α → β) (g : β → γ) (xs : List α) :
    versionA f g xs = versionB f g xs := by
  exact Specification.unique (versionA_correct f g xs) (versionB_correct f g xs)

Definition 3 (Traversal cost).

Assign one unit of cost to each nonempty list cell visited by a map:

countedMap f xs returns a pair $(\text{output},\text{cost})$. Its second component implements this recurrence: zero for [], and cost + 1 for x :: xs.

def countedMap (f : α → β) : List α → List β × Nat
  | [] => ([], 0)
  | x :: xs =>
    let (ys, cost) := countedMap f xs
    (f x :: ys, cost + 1)

The counted executions versionACounted and versionBCounted follow the two implementations. A adds the costs of its two maps; B counts one composed map.

def versionACounted (f : α → β) (g : β → γ) (xs : List α) : List γ × Nat :=
  let (ys, firstCost) := countedMap f xs
  let (zs, secondCost) := countedMap g ys
  (zs, firstCost + secondCost)

def versionBCounted (f : α → β) (g : β → γ) (xs : List α) : List γ × Nat :=
  countedMap (fun x => g (f x)) xs

Write $\pi_2$ for the second component of a pair (.2 in Lean). The costs below refer precisely to these executions:

Lemma 4 (One map).

A counted map returns the ordinary map result and visits exactly $|x|$ cells:

Proof.

Induction on $x$: the empty case has cost zero; the nonempty case adds one visit. This is countedMap_spec:

@[simp] theorem countedMap_spec (f : α → β) (xs : List α) :
    countedMap f xs = (xs.map f, xs.length) := by
  induction xs with
  | nil => rfl
  | cons x xs ih => simp [countedMap, ih]

Theorem 5 (Exact traversal costs).

The counted executions preserve the original outputs and record their costs:

In particular, $C_A(x)=2|x|$ and $C_B(x)=|x|$.

Proof.

Apply countedMap_spec twice for A and once for B. Mapping preserves list length. The two conjuncts of implementation_costs are exactly the two equalities above:

theorem implementation_costs (f : α → β) (g : β → γ) (xs : List α) :
    versionACounted f g xs = (versionA f g xs, 2 * xs.length) ∧
    versionBCounted f g xs = (versionB f g xs, xs.length) := by
  simp [versionACounted, versionBCounted, versionA, versionB, Nat.two_mul]