1	# A pickle, the object graph torch.save writes, replayed opcode by opcode into
2	# values: the stack machine of protocol 2 with the string, global and memo
3	# opcodes of protocol 4. No code runs: a class is a Global, an instance keeps
4	# its class, arguments and state, and a call keeps its callable and arguments.
5	import Base
6	import ./storage_bytes.bend as Bytes
7	import ./storage_buffer.bend as Buffers
8	
9	type Value is Data:
10	  Absent{}
11	  Flag{value: Bool}
12	  Integer{word: U32}
13	  Float{high: U32,low: U32}
14	  Text{text: String}
15	  Sequence{items: List<&2,Value>}
16	  Items{items: List<&2,Value>}
17	  Dict{entries: List<&2,Value>}
18	  Global{module: String,name: String}
19	  Object{class: Value,arguments: Value,state: Value}
20	  Reduced{callable: Value,arguments: Value}
21	  Persistent{id: Value}
22	  Mark{}
23	
24	type Parsed is Data:
25	  Parsed{value: Value}
26	  Rejected{detail: String}
27	
28	# The stack, the memo with its capacity, and the index the next memoize takes.
29	type Machine is Type:
30	  Machine{stack: List<&2,Value>,memo: Array<Value>,capacity: U32,next: U32}
31	
32	type Step is Type:
33	  Running{bytes: List<&2,U32>,machine: Machine}
34	  Completed{value: Value}
35	  Halted{detail: String}
36	
37	type Split is Data:
38	  Split{items: List<&2,Value>,rest: List<&2,Value>}
39	  Unmarked{}
40	
41	type Line is Data:
42	  Line{text: List<&2,U32>,rest: List<&2,U32>}
43	
44	# Dictionaries: entries are two-item sequences; assignment replaces an existing key.
45	
46	def same_key(left: Value,right: Value) -> Bool:
47	  match left right:
48	    case Text{left} Text{right}: String.eq(left,right)
49	    case Integer{left} Integer{right}: U32.is_eq(left,right)
50	    case left right: False{}
51	
52	def key_of(entry: Value) -> Value:
53	  match entry:
54	    case Sequence{Con{key,others}}: key
55	    case entry: Absent{}
56	
57	def assign(entries: List<&2,Value>,+key: Value,+value: Value) -> List<&2,Value>:
58	  match entries:
59	    case Nil{}: [Sequence{[key,value]}]
60	    case Con{+entry,+rest}: Bool.pick(List<&2,Value>,same_key(key_of(entry),key),Sequence{[key,value]} <> rest,entry <> assign(rest,key,value))
61	
62	def assign_pairs(pairs: List<&2,Value>,entries: List<&2,Value>) -> List<&2,Value>:
63	  match pairs:
64	    case Con{key,Con{value,more}}: assign_pairs(more,assign(entries,key,value))
65	    case pairs: entries
66	
67	# The value under a text key, or Absent.
68	def get(entries: List<&2,Value>,+key: String) -> Value:
69	  match entries:
70	    case Nil{}: Absent{}
71	    case Con{Sequence{Con{Text{+name},Con{+value,others}}},rest}: Bool.pick(Value,String.eq(name,key),value,get(rest,key))
72	    case Con{entry,rest}: get(rest,key)
73	
74	# The items above the mark, oldest first, and the stack below it.
75	def split(stack: List<&2,Value>,items: List<&2,Value>) -> Split:
76	  match stack:
77	    case Nil{}: Unmarked{}
78	    case Con{Mark{},rest}: Split{items,rest}
79	    case Con{top,rest}: split(rest,top <> items)
80	
81	def line(bytes: List<&2,U32>,ended: Bool,+text: List<&2,U32>) -> Line:
82	  match bytes ended:
83	    case bytes True{}: Line{List.reverse(&2,U32,text),bytes}
84	    case Nil{} False{}: Line{List.reverse(&2,U32,text),[]}
85	    case Con{+byte,rest} False{}: line(rest,U32.is_eq(byte,10),Bool.pick(List<&2,U32>,U32.is_eq(byte,10),text,byte <> text))
86	
87	# Little-endian integers of one to four bytes, sign-extended as the LONG1 opcode means them.
88	def sign_extended(+length: U32,+word: U32) -> Value:
89	  match length:
90	    case 1: Integer{Bool.pick(U32,U32.is_le(128,word),(word + 4294967040 : U32),word)}
91	    case 2: Integer{Bool.pick(U32,U32.is_le(32768,word),(word + 4294901760 : U32),word)}
92	    case 3: Integer{Bool.pick(U32,U32.is_le(8388608,word),(word + 4278190080 : U32),word)}
93	    case _: Integer{word}
94	
95	def little_endian(bytes: List<&2,U32>,+weight: U32) -> U32:
96	  match bytes:
97	    case Nil{}: 0
98	    case Con{byte,rest}: (byte * weight + little_endian(rest,(weight * 256 : U32)) : U32)
99	
100	# Opcodes
101	
102	def push(value: Value,machine: Machine) -> Machine:
103	  Machine{stack,memo,capacity,next} = machine
104	  Machine{value <> stack,memo,capacity,next}
105	
106	def pushed(bytes: List<&2,U32>,value: Value,machine: Machine) -> Step:
107	  Running{bytes,push(value,machine)}
108	
109	def stored(bytes: List<&2,U32>,fits: Bool,+index: U32,+top: Value,stack: List<&2,Value>,memo: Array<Value>,+capacity: U32,+next: U32) -> Step:
110	  match fits:
111	    case False{}: Halted{"memo index " ++ U32.show(index) ++ " exceeds the pickle"}
112	    case True{}: Running{bytes,Machine{top <> stack,Array.set(Value,memo,index,top),capacity,U32.max(next,(index + 1 : U32))}}
113	
114	# BINPUT, LONG_BINPUT and MEMOIZE keep the top of the stack under an index.
115	def remember(bytes: List<&2,U32>,+index: U32,machine: Machine) -> Step:
116	  Machine{stack,memo,+capacity,+next} = machine
117	  match stack:
118	    case Nil{}: Halted{"PUT on an empty stack"}
119	    case Con{+top,rest}: stored(bytes,U32.is_lt(index,capacity),index,top,rest,memo,capacity,next)
120	
121	def recalled(bytes: List<&2,U32>,pair: Array<Value> & Value,stack: List<&2,Value>,+capacity: U32,+next: U32) -> Step:
122	  (memo,value) = pair
123	  Running{bytes,Machine{value <> stack,memo,capacity,next}}
124	
125	def recall_within(bytes: List<&2,U32>,fits: Bool,+index: U32,stack: List<&2,Value>,memo: Array<Value>,+capacity: U32,+next: U32) -> Step:
126	  match fits:
127	    case False{}: Halted{"memo index " ++ U32.show(index) ++ " exceeds the pickle"}
128	    case True{}: recalled(bytes,Array.get(Value,memo,index),stack,capacity,next)
129	
130	# BINGET and LONG_BINGET push the value under an index.
131	def recall(bytes: List<&2,U32>,+index: U32,machine: Machine) -> Step:
132	  Machine{stack,memo,+capacity,+next} = machine
133	  recall_within(bytes,U32.is_lt(index,capacity),index,stack,memo,capacity,next)
134	
135	def memoize(bytes: List<&2,U32>,machine: Machine) -> Step:
136	  Machine{stack,memo,+capacity,+next} = machine
137	  remember(bytes,next,Machine{stack,memo,capacity,next})
138	
139	def put_short(bytes: List<&2,U32>,machine: Machine) -> Step:
140	  match bytes:
141	    case Con{index,rest}: remember(rest,index,machine)
142	    case bytes: Halted{"BINPUT without its index"}
143	
144	def put_long(bytes: List<&2,U32>,machine: Machine) -> Step:
145	  match bytes:
146	    case Con{first,Con{second,Con{third,Con{fourth,rest}}}}: remember(rest,little_endian([first,second,third,fourth],1),machine)
147	    case bytes: Halted{"LONG_BINPUT without its index"}
148	
149	def get_short(bytes: List<&2,U32>,machine: Machine) -> Step:
150	  match bytes:
151	    case Con{index,rest}: recall(rest,index,machine)
152	    case bytes: Halted{"BINGET without its index"}
153	
154	def get_long(bytes: List<&2,U32>,machine: Machine) -> Step:
155	  match bytes:
156	    case Con{first,Con{second,Con{third,Con{fourth,rest}}}}: recall(rest,little_endian([first,second,third,fourth],1),machine)
157	    case bytes: Halted{"LONG_BINGET without its index"}
158	
159	# Text: a length, then that many UTF-8 bytes.
160	def text_of(+length: U32,+bytes: List<&2,U32>,machine: Machine) -> Step:
161	  Running{List.drop(&2,U32,bytes,U32.to_nat(length)),push(Text{Bytes.text(List.take(&2,U32,bytes,U32.to_nat(length)))},machine)}
162	
163	def text_short(bytes: List<&2,U32>,machine: Machine) -> Step:
164	  match bytes:
165	    case Con{length,rest}: text_of(length,rest,machine)
166	    case bytes: Halted{"a short string without its length"}
167	
168	def text_long(bytes: List<&2,U32>,machine: Machine) -> Step:
169	  match bytes:
170	    case Con{first,Con{second,Con{third,Con{fourth,rest}}}}: text_of(little_endian([first,second,third,fourth],1),rest,machine)
171	    case bytes: Halted{"a string without its length"}
172	
173	# GLOBAL: a module and a name, each ended by a newline.
174	def global_named(second: Line,+module: List<&2,U32>,machine: Machine) -> Step:
175	  Line{name,rest} = second
176	  Running{rest,push(Global{Bytes.text(module),Bytes.text(name)},machine)}
177	
178	def global_module(first: Line,machine: Machine) -> Step:
179	  Line{module,rest} = first
180	  global_named(line(rest,False{},[]),module,machine)
181	
182	def global(bytes: List<&2,U32>,machine: Machine) -> Step:
183	  global_module(line(bytes,False{},[]),machine)
184	
185	# STACK_GLOBAL takes the name and the module from the stack.
186	def stack_global(bytes: List<&2,U32>,machine: Machine) -> Step:
187	  Machine{stack,memo,capacity,next} = machine
188	  match stack:
189	    case Con{Text{name},Con{Text{module},rest}}: Running{bytes,Machine{Global{module,name} <> rest,memo,capacity,next}}
190	    case stack: Halted{"STACK_GLOBAL without two strings"}
191	
192	# Numbers
193	
194	def integer_short(bytes: List<&2,U32>,machine: Machine) -> Step:
195	  match bytes:
196	    case Con{byte,rest}: pushed(rest,Integer{byte},machine)
197	    case bytes: Halted{"BININT1 without its byte"}
198	
199	def integer_two(bytes: List<&2,U32>,machine: Machine) -> Step:
200	  match bytes:
201	    case Con{first,Con{second,rest}}: pushed(rest,Integer{little_endian([first,second],1)},machine)
202	    case bytes: Halted{"BININT2 without its bytes"}
203	
204	def integer_four(bytes: List<&2,U32>,machine: Machine) -> Step:
205	  match bytes:
206	    case Con{first,Con{second,Con{third,Con{fourth,rest}}}}: pushed(rest,Integer{little_endian([first,second,third,fourth],1)},machine)
207	    case bytes: Halted{"BININT without its bytes"}
208	
209	def long_sized(narrow: Bool,+length: U32,+bytes: List<&2,U32>,machine: Machine) -> Step:
210	  match narrow:
211	    case False{}: Halted{"an integer wider than 32 bits"}
212	    case True{}: Running{List.drop(&2,U32,bytes,U32.to_nat(length)),push(sign_extended(length,little_endian(List.take(&2,U32,bytes,U32.to_nat(length)),1)),machine)}
213	
214	# LONG1: a byte count, then the two's complement bytes.
215	def long_one(bytes: List<&2,U32>,machine: Machine) -> Step:
216	  match bytes:
217	    case Con{+length,rest}: long_sized(U32.is_le(length,4),length,rest,machine)
218	    case bytes: Halted{"LONG1 without its length"}
219	
220	# BINFLOAT: eight bytes, most significant first.
221	def float(bytes: List<&2,U32>,machine: Machine) -> Step:
222	  match bytes:
223	    case Con{a,Con{b,Con{c,Con{d,Con{e,Con{f,Con{g,Con{h,rest}}}}}}}}:
224	      pushed(rest,Float{little_endian([d,c,b,a],1),little_endian([h,g,f,e],1)},machine)
225	    case bytes: Halted{"BINFLOAT without its bytes"}
226	
227	# Containers
228	
229	def tuple_of(parts: Split,bytes: List<&2,U32>,memo: Array<Value>,+capacity: U32,+next: U32) -> Step:
230	  match parts:
231	    case Unmarked{}: Halted{"TUPLE without a mark"}
232	    case Split{items,rest}: Running{bytes,Machine{Sequence{items} <> rest,memo,capacity,next}}
233	
234	def tuple(bytes: List<&2,U32>,machine: Machine) -> Step:
235	  Machine{stack,memo,+capacity,+next} = machine
236	  tuple_of(split(stack,[]),bytes,memo,capacity,next)
237	
238	def tuple_one(bytes: List<&2,U32>,machine: Machine) -> Step:
239	  Machine{stack,memo,capacity,next} = machine
240	  match stack:
241	    case Con{first,rest}: Running{bytes,Machine{Sequence{[first]} <> rest,memo,capacity,next}}
242	    case stack: Halted{"TUPLE1 on an empty stack"}
243	
244	def tuple_two(bytes: List<&2,U32>,machine: Machine) -> Step:
245	  Machine{stack,memo,capacity,next} = machine
246	  match stack:
247	    case Con{second,Con{first,rest}}: Running{bytes,Machine{Sequence{[first,second]} <> rest,memo,capacity,next}}
248	    case stack: Halted{"TUPLE2 on a short stack"}
249	
250	def tuple_three(bytes: List<&2,U32>,machine: Machine) -> Step:
251	  Machine{stack,memo,capacity,next} = machine
252	  match stack:
253	    case Con{third,Con{second,Con{first,rest}}}: Running{bytes,Machine{Sequence{[first,second,third]} <> rest,memo,capacity,next}}
254	    case stack: Halted{"TUPLE3 on a short stack"}
255	
256	def items_appended(parts: Split,bytes: List<&2,U32>,memo: Array<Value>,+capacity: U32,+next: U32) -> Step:
257	  match parts:
258	    case Unmarked{}: Halted{"APPENDS without a mark"}
259	    case Split{items,Con{Items{existing},rest}}: Running{bytes,Machine{Items{List.append(&2,Value,existing,items)} <> rest,memo,capacity,next}}
260	    case Split{items,rest}: Halted{"APPENDS on a value that is not a list"}
261	
262	def appends(bytes: List<&2,U32>,machine: Machine) -> Step:
263	  Machine{stack,memo,+capacity,+next} = machine
264	  items_appended(split(stack,[]),bytes,memo,capacity,next)
265	
266	def append(bytes: List<&2,U32>,machine: Machine) -> Step:
267	  Machine{stack,memo,capacity,next} = machine
268	  match stack:
269	    case Con{value,Con{Items{existing},rest}}: Running{bytes,Machine{Items{List.append(&2,Value,existing,[value])} <> rest,memo,capacity,next}}
270	    case stack: Halted{"APPEND on a value that is not a list"}
271	
272	def entries_assigned(parts: Split,bytes: List<&2,U32>,memo: Array<Value>,+capacity: U32,+next: U32) -> Step:
273	  match parts:
274	    case Unmarked{}: Halted{"SETITEMS without a mark"}
275	    case Split{pairs,Con{Dict{entries},rest}}: Running{bytes,Machine{Dict{assign_pairs(pairs,entries)} <> rest,memo,capacity,next}}
276	    case Split{pairs,rest}: Halted{"SETITEMS on a value that is not a dictionary"}
277	
278	def set_items(bytes: List<&2,U32>,machine: Machine) -> Step:
279	  Machine{stack,memo,+capacity,+next} = machine
280	  entries_assigned(split(stack,[]),bytes,memo,capacity,next)
281	
282	def set_item(bytes: List<&2,U32>,machine: Machine) -> Step:
283	  Machine{stack,memo,capacity,next} = machine
284	  match stack:
285	    case Con{value,Con{key,Con{Dict{entries},rest}}}: Running{bytes,Machine{Dict{assign(entries,key,value)} <> rest,memo,capacity,next}}
286	    case stack: Halted{"SETITEM on a value that is not a dictionary"}
287	
288	# Objects and calls
289	
290	def new_object(bytes: List<&2,U32>,machine: Machine) -> Step:
291	  Machine{stack,memo,capacity,next} = machine
292	  match stack:
293	    case Con{arguments,Con{class,rest}}: Running{bytes,Machine{Object{class,arguments,Absent{}} <> rest,memo,capacity,next}}
294	    case stack: Halted{"NEWOBJ on a short stack"}
295	
296	def build(bytes: List<&2,U32>,machine: Machine) -> Step:
297	  Machine{stack,memo,capacity,next} = machine
298	  match stack:
299	    case Con{state,Con{Object{class,arguments,old},rest}}: Running{bytes,Machine{Object{class,arguments,state} <> rest,memo,capacity,next}}
300	    case stack: Halted{"BUILD on a value that is not an object"}
301	
302	# An OrderedDict call with no arguments is a dictionary, so its items can follow.
303	def called(ordered: Bool,callable: Value,arguments: Value) -> Value:
304	  match ordered:
305	    case True{}: Dict{[]}
306	    case False{}: Reduced{callable,arguments}
307	
308	def reduced(callable: Value,arguments: Value) -> Value:
309	  match callable arguments:
310	    case Global{+module,+name} Sequence{Nil{}}:
311	      called(String.eq(module,"collections") && String.eq(name,"OrderedDict"),Global{module,name},Sequence{[]})
312	    case callable arguments: Reduced{callable,arguments}
313	
314	def reduce(bytes: List<&2,U32>,machine: Machine) -> Step:
315	  Machine{stack,memo,capacity,next} = machine
316	  match stack:
317	    case Con{arguments,Con{callable,rest}}: Running{bytes,Machine{reduced(callable,arguments) <> rest,memo,capacity,next}}
318	    case stack: Halted{"REDUCE on a short stack"}
319	
320	def persistent(bytes: List<&2,U32>,machine: Machine) -> Step:
321	  Machine{stack,memo,capacity,next} = machine
322	  match stack:
323	    case Con{id,rest}: Running{bytes,Machine{Persistent{id} <> rest,memo,capacity,next}}
324	    case stack: Halted{"BINPERSID on an empty stack"}
325	
326	def stop(machine: Machine) -> Step:
327	  Machine{stack,memo,capacity,next} = machine
328	  match stack:
329	    case Con{value,rest}: Completed{value}
330	    case stack: Halted{"STOP on an empty stack"}
331	
332	def skip(count: Nat,bytes: List<&2,U32>,machine: Machine) -> Step:
333	  Running{List.drop(&2,U32,bytes,count),machine}
334	
335	# One opcode and its argument bytes.
336	def operate(+code: U32,bytes: List<&2,U32>,machine: Machine) -> Step:
337	  match code:
338	    case 128: skip(1n,bytes,machine)
339	    case 149: skip(8n,bytes,machine)
340	    case 46: stop(machine)
341	    case 40: pushed(bytes,Mark{},machine)
342	    case 78: pushed(bytes,Absent{},machine)
343	    case 136: pushed(bytes,Flag{True{}},machine)
344	    case 137: pushed(bytes,Flag{False{}},machine)
345	    case 41: pushed(bytes,Sequence{[]},machine)
346	    case 93: pushed(bytes,Items{[]},machine)
347	    case 125: pushed(bytes,Dict{[]},machine)
348	    case 113: put_short(bytes,machine)
349	    case 114: put_long(bytes,machine)
350	    case 104: get_short(bytes,machine)
351	    case 106: get_long(bytes,machine)
352	    case 148: memoize(bytes,machine)
353	    case 88: text_long(bytes,machine)
354	    case 140: text_short(bytes,machine)
355	    case 84: text_long(bytes,machine)
356	    case 85: text_short(bytes,machine)
357	    case 99: global(bytes,machine)
358	    case 147: stack_global(bytes,machine)
359	    case 75: integer_short(bytes,machine)
360	    case 77: integer_two(bytes,machine)
361	    case 74: integer_four(bytes,machine)
362	    case 138: long_one(bytes,machine)
363	    case 71: float(bytes,machine)
364	    case 116: tuple(bytes,machine)
365	    case 133: tuple_one(bytes,machine)
366	    case 134: tuple_two(bytes,machine)
367	    case 135: tuple_three(bytes,machine)
368	    case 101: appends(bytes,machine)
369	    case 97: append(bytes,machine)
370	    case 117: set_items(bytes,machine)
371	    case 115: set_item(bytes,machine)
372	    case 129: new_object(bytes,machine)
373	    case 98: build(bytes,machine)
374	    case 82: reduce(bytes,machine)
375	    case 81: persistent(bytes,machine)
376	    case _: Halted{"opcode " ++ U32.show(code) ++ " is not supported"}
377	
378	def step(bytes: List<&2,U32>,machine: Machine) -> Step:
379	  match bytes:
380	    case Nil{}: Halted{"ends before STOP"}
381	    case Con{code,rest}: operate(code,rest,machine)
382	
383	# Every opcode takes at least one byte, so the byte count bounds the steps.
384	def run(fuel: Nat,state: Step) -> Parsed:
385	  match fuel state:
386	    case fuel Completed{value}: Parsed{value}
387	    case fuel Halted{detail}: Rejected{detail}
388	    case 0n Running{bytes,machine}: Rejected{"does not end"}
389	    case 1n+rest Running{bytes,machine}: run(rest,step(bytes,machine))
390	
391	# A memo index is below the count of PUT opcodes, each of two bytes or more.
392	def parse(+bytes: List<&2,U32>) -> Parsed:
393	  +count = List.length(&2,U32,bytes)
394	  +depth = Buffers.depth_for((U32.from_nat(count) / 2 + 1 : U32))
395	  run(count,Running{bytes,Machine{[],Array.new(Value,depth,Absent{}),U32.shln(1,depth),0}})