leanexe/Verified/Reflect/Modes.lean at main · jsmorph/leanexe · GitHub

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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
import Verified.Compile

/-! The modes that the reflector chooses for a function's parameters.  The choice is not trusted:
the correctness theorem holds for every choice of modes.  An array parameter is owned when the
body consumes it where it dies, where a borrowed parameter would be copied: as the array of an
update or an extension, at a call that moves it into an owned parameter, as the function's
result, or through a binding, a branch, a pair, or a loop's state into one of these.  Every other
parameter is borrowed, so that a caller need not copy an array that the function only reads.
The live sets follow those of `Expr.code`. -/

namespace Verified

/-- Whether an expression consumes variable `k` where it dies, given the variables `live` live
after the expression and whether its value is consumed, `owned`.  A value is consumed when it
is the function's result, a value that a binding's body consumes, a branch or a component of a
value that is consumed or owned, or a loop's state when its body consumes the state or returns
an owned value. -/
def Expr.demands : {Γ : List Ty} → {t : Ty} → Expr S Γ t → (Nat → Bool) → Bool → Nat → Bool
  | _, _, .var x, live, owned, k => owned && k == x.index && !live k
  | _, _, .word _, _, _, _ | _, _, .bool _, _, _, _ | _, _, .size _, _, _, _
  | _, _, .leSize _ _, _, _, _
  | _, _, .float _, _, _, _ | _, _, .float32 _, _, _, _ => false
  | _, _, .treeNil _ _, _, _, _ => false
  | _, _, .bytesEmpty, _, _, _ | _, _, .bytesLiteral _, _, _, _
  | _, _, .bytesSize _, _, _, _ => false
  | _, _, .bytesGet x index, live, _, k =>
    index.demands (fun j => live j || j == x.index) false k
  | _, _, .treeNode _ reuse words arrays buffers children, live, _, k =>
    let kept := fun i => live i || optionVarUses reuse i
    (optionVarUses reuse k && !live k) ||
      (argsAny fun i => (words i).demands
        (fun j => kept j || (argsAny fun next => decide (i.val < next.val) && (words next).uses j) ||
          (argsAny fun next => (children next).uses j) ||
          (argsAny fun next => (arrays next).uses j) ||
          (argsAny fun next => (buffers next).uses j)) false k) ||
      (argsAny fun i => (children i).demands
        (fun j => kept j ||
          (argsAny fun next => decide (i.val < next.val) && (children next).uses j) ||
          (argsAny fun next => (arrays next).uses j) ||
          (argsAny fun next => (buffers next).uses j)) true k) ||
      (argsAny fun i => (arrays i).demands
        (fun j => kept j ||
          (argsAny fun next => decide (i.val < next.val) && (arrays next).uses j) ||
          (argsAny fun next => (buffers next).uses j)) true k) ||
      (argsAny fun i => (buffers i).demands
        (fun j => kept j ||
          (argsAny fun next => decide (i.val < next.val) && (buffers next).uses j)) true k)
  | _, _, .treeCase (shape := shape) x empty nodes, live, owned, k =>
    let branch := owned ||
      (if h : shape.hasNull = true then (empty h).mode [] == .owned else false) ||
      argsAny fun r => (nodes r).mode [] == .owned
    let nodeDemands := fun r : Fin shape.records =>
      let inner := shift (shape.bindCount r) live
      let fields := (nodes r).demands inner branch 0 ||
        (argsAny fun i : Fin (shape.children r) =>
          (nodes r).demands inner branch (1 + shape.words r + i.val))
      (k == x.index && !live k && !(nodes r).uses (k + shape.bindCount r) && fields) ||
        (nodes r).demands inner branch (k + shape.bindCount r)
    (if h : shape.hasNull = true then (empty h).demands live branch k else false) ||
      argsAny nodeDemands
  | _, _, .bin _ left right, live, _, k | _, _, .cmp _ left right, live, _, k
  | _, _, .and left right, live, _, k | _, _, .or left right, live, _, k
  | _, _, .fbin _ left right, live, _, k | _, _, .fcmp _ left right, live, _, k
  | _, _, .fbin32 _ left right, live, _, k | _, _, .fcmp32 _ left right, live, _, k =>
    left.demands (fun i => live i || right.uses i) false k || right.demands live false k
  | _, _, .not e, live, _, k | _, _, .funary _ e, live, _, k => e.demands live false k
  | _, _, .toFloat _ e, live, _, k => e.demands live false k
  | _, _, .toWord _ e, live, _, k | _, _, .funary32 _ e, live, _, k
  | _, _, .toFloat32 _ e, live, _, k | _, _, .toWord32 _ e, live, _, k => e.demands live false k
  | _, _, .mk first second, live, _, k =>
    first.demands (fun i => live i || second.uses i) false k || second.demands live false k
  | _, _, .proj _ _, _, _, _ => false
  | _, _, .ite c thenE elseE, live, owned, k =>
    let branch := owned || thenE.mode [] == .owned || elseE.mode [] == .owned
    c.demands (fun i => live i || thenE.uses i || elseE.uses i) false k ||
      thenE.demands live branch k || elseE.demands live branch k
  | _, _, .letE value body, live, owned, k =>
    value.demands (fun i => live i || body.uses (i + 1)) (body.demands (shift 1 live) owned 0) k ||
      body.demands (shift 1 live) owned (k + 1)
  | _, _, .call (g := g) _ args, live, _, k =>
    argsAny fun i =>
      if (g.params.get i).scalar then
        (args i).demands (fun j => live j || argsAny fun j' => (args j').uses j) false k
      else
        g.mode i == .owned && (args i).varIndex? == some k && !live k &&
          !(argsAny fun j => j != i && !(g.params.get j).scalar && (args j).uses k)
  | _, _, .pair first second, live, owned, k =>
    let component := owned || first.mode [] == .owned || second.mode [] == .owned
    first.demands (fun i => live i || second.uses i) component k ||
      second.demands live component k
  | _, _, .letPair e body, live, owned, k =>
    let inner := shift 2 live
    e.demands (fun i => live i || body.uses (i + 2))
        (body.demands inner owned 0 || body.demands inner owned 1) k ||
      body.demands inner owned (k + 2)
  | _, _, .loop count init cond body, live, owned, k =>
    let all := fun i => live i || cond.uses (i + 1) || body.uses (i + 2)
    let inner := shift 2 all
    let state := owned || body.demands inner false 0 || body.mode [] == .owned
    count.demands (fun i => all i || init.uses i) false k || init.demands all state k ||
      cond.demands (fun j => j == 0 || shift 1 all j) false (k + 1) ||
      body.demands inner state (k + 2)
  | _, _, .fold x init body, live, owned, k =>
    let all := fun i => live i || (Expr.foldStep x body).uses (i + 2)
    let inner := shift 1 (shift 2 all)
    let state := owned || body.demands inner false 1 || body.mode [.borrowed] == .owned
    init.demands all state k || body.demands inner state (k + 3)
  | _, _, .listFold (e := e) (captures := captures) args _, live, _, k =>
    argsAny fun i =>
      if ((ListFold.params e captures).get i).scalar then
        (args i).demands (fun j => live j || argsAny fun j' => (args j').uses j) false k
      else false
  | _, _, .get x i, live, _, k => i.demands (fun j => live j || j == x.index) false k
  | _, _, .build count elem, live, _, k | _, _, .bytesBuild count elem, live, _, k =>
    let all := fun i => live i || elem.uses (i + 1)
    count.demands all false k || elem.demands (shift 1 all) false (k + 1)
  | _, _, .set x i v, live, _, k =>
    (k == x.index && !live k) || i.demands (fun j => live j || v.uses j || j == x.index) false k ||
      v.demands (fun j => live j || j == x.index) false k
  | _, _, .push x v, live, _, k =>
    (k == x.index && !live k) || v.demands (fun j => live j || j == x.index) false k
  | _, _, .append x y, live, _, k => k == x.index && !live k && x.index != y.index
  | _, _, .insertAt x i v, live, _, k =>
    (k == x.index && !live k) || i.demands (fun j => live j || v.uses j || j == x.index) false k ||
      v.demands (fun j => live j || j == x.index) false k
  | _, _, .eraseAt x i, live, _, k =>
    (k == x.index && !live k) || i.demands (fun j => live j || j == x.index) false k

/-- The modes chosen for the parameters of types `params` of a function with body `body`: owned
for an array parameter that the body consumes where it dies, and borrowed otherwise. -/
def Expr.paramChoice {params : List Ty} {result : Ty} (body : Expr S params result) :
    List Mode :=
  (List.range params.length).map fun i =>
    match params[i]? with
    | some (Ty.array _) | some (Ty.tree _) | some (Ty.treeHeader _ _) | some Ty.bytes =>
      if body.demands (fun _ => false) true i then .owned else .borrowed
    | _ => .borrowed

end Verified