1	# Raw little-endian FP32 files and NumPy .npy files through the official File effects.
2	import Base
3	import ./storage_buffer.bend as Buffers
4	import ./tensor.bend as Tensors
5	import ./tensor_layout.bend as Layout
6	
7	def Tensor() -> Type:
8	  Tensors.Tensor<F32>
9	
10	def Read() -> Type:
11	  File & Result<&1,&1,U32 & String,List<&2,U32>>
12	
13	def Written() -> Type:
14	  File & Result<&1,&1,U32 & String,Unit>
15	
16	# Values per File effect.
17	def chunk() -> U32:
18	  16384
19	
20	# An FP32 value and its 32 bits are the same word.
21	
22	def from_bits(word: U32) -> F32:
23	  U32{bits} = word
24	  F32{bits}
25	
26	def to_bits(value: F32) -> U32:
27	  F32{bits} = value
28	  U32{bits}
29	
30	# Reading: bytes gather into words, least significant first. The values of one
31	# read are written into the buffer before the next read, so no list outgrows a read.
32	
33	type Words is Data:
34	  Words{values: List<&2,F32>,count: U32,word: U32,weight: U32}
35	
36	# The buffer of the tensor, the values found so far, and a word in progress.
37	type Target is Type:
38	  Target{buffer: Buffers.Buffer<F32>,found: U32,word: U32,weight: U32}
39	
40	type Reading is Type:
41	  More{file: File,target: Target}
42	  End{file: File,target: Target,complete: Bool}
43	
44	type Outcome is Type:
45	  Outcome{target: Target,complete: Bool}
46	
47	def gather(full: Bool,values: List<&2,F32>,+count: U32,+word: U32,+weight: U32) -> Words:
48	  match full:
49	    case False{}: Words{values,count,word,(weight * 256 : U32)}
50	    case True{}: Words{from_bits(word) <> values,(count + 1 : U32),0,1}
51	
52	def push(+byte: U32,words: Words) -> Words:
53	  Words{values,+count,+word,+weight} = words
54	  gather(U32.is_eq(weight,16777216),values,count,(word + byte * weight : U32),weight)
55	
56	def trickle(bytes: List<&2,U32>,words: Words) -> Words:
57	  match bytes:
58	    case []: words
59	    case byte <> rest: trickle(rest,push(byte,words))
60	
61	# Four bytes at a time while the bytes are whole words; the values are newest first.
62	def quads(bytes: List<&2,U32>,values: List<&2,F32>,+count: U32) -> Words:
63	  match bytes:
64	    case Con{first,Con{second,Con{third,Con{fourth,rest}}}}:
65	      quads(rest,from_bits((first + second * 256 + third * 65536 + fourth * 16777216 : U32)) <> values,(count + 1 : U32))
66	    case rest: trickle(rest,Words{values,count,0,1})
67	
68	def words(whole: Bool,bytes: List<&2,U32>,+word: U32,+weight: U32) -> Words:
69	  match whole:
70	    case True{}: quads(bytes,[],0)
71	    case False{}: trickle(bytes,Words{[],0,word,weight})
72	
73	# Values beyond the count of the shape are counted and not written.
74	def placed(read: Words,buffer: Buffers.Buffer<F32>,+found: U32,+count: U32) -> Target:
75	  Words{values,+added,+word,+weight} = read
76	  Target{Buffers.write_at(~F32,List.take(&2,F32,List.reverse(&2,F32,values),U32.to_nat((count - U32.min(found,count) : U32))),U32.min(found,count),buffer),
77	    (found + added : U32),word,weight}
78	
79	def fill(bytes: List<&2,U32>,target: Target,+count: U32) -> Target:
80	  Target{buffer,+found,+word,+weight} = target
81	  placed(words(U32.is_eq(weight,1),bytes,word,weight),buffer,found,count)
82	
83	def absorbed(bytes: List<&2,U32>,file: File,target: Target,+count: U32) -> Reading:
84	  match bytes:
85	    case []: End{file,target,True{}}
86	    case byte <> rest: More{file,fill(byte <> rest,target,count)}
87	
88	def received(result: Result<&1,&1,U32 & String,List<&2,U32>>,file: File,target: Target,+count: U32) -> Reading:
89	  match result:
90	    case Fail{error}: End{file,target,False{}}
91	    case Done{bytes}: absorbed(bytes,file,target,count)
92	
93	def taken(read: Read(),target: Target,+count: U32) -> Reading:
94	  (file,result) = read
95	  received(result,file,target,count)
96	
97	def closed(file: File,target: Target,complete: Bool) -> IO(Outcome):
98	  do IO<Outcome>:
99	    File.close(file)
100	    IO.pure(Outcome,Outcome{target,complete})
101	
102	# A file longer than its shape stops at the fuel and fails the count.
103	def read_all(fuel: Nat,reading: Reading,+count: U32) -> IO(Outcome):
104	  match fuel reading:
105	    case 0n More{file,target}: closed(file,target,True{})
106	    case 0n End{file,target,complete}: closed(file,target,complete)
107	    case 1n+rest End{file,target,complete}: closed(file,target,complete)
108	    case 1n+rest More{file,target}:
109	      do IO<Outcome>:
110	        read : Read() <- File.read_bytes(file,(chunk() * 4 : U32))
111	        read_all(rest,taken(read,target,count),count)
112	
113	def fresh(+count: U32) -> Target:
114	  Target{Tensors.buffer(~F32,0.0,count),0,0,1}
115	
116	def fuel(+count: U32) -> Nat:
117	  U32.to_nat((count / chunk() + 2 : U32))
118	
119	def fault(operation: String,path: String,detail: String,shape: List<&2,U32>) -> Tensor():
120	  Tensors.failed(~F32,operation,Tensors.FileFault{path,detail},shape,0.0,1)
121	
122	def sized(matches: Bool,operation: String,path: String,+shape: List<&2,U32>,buffer: Buffers.Buffer<F32>) -> Tensor():
123	  match matches:
124	    case True{}: Tensors.from_storage(~F32,shape,buffer,0.0)
125	    case False{}: fault(operation,path,"does not hold the values of the shape",shape)
126	
127	def whole(aligned: Bool,operation: String,path: String,+shape: List<&2,U32>,buffer: Buffers.Buffer<F32>,+found: U32,+count: U32) -> Tensor():
128	  match aligned:
129	    case False{}: fault(operation,path,"is not a whole number of FP32 values",shape)
130	    case True{}: sized(U32.is_eq(found,count),operation,path,shape,buffer)
131	
132	def aligned(operation: String,path: String,+shape: List<&2,U32>,target: Target,+count: U32) -> Tensor():
133	  Target{buffer,+found,word,+weight} = target
134	  whole(U32.is_eq(weight,1),operation,path,shape,buffer,found,count)
135	
136	def checked(complete: Bool,operation: String,path: String,+shape: List<&2,U32>,target: Target,+count: U32) -> Tensor():
137	  match complete:
138	    case False{}: fault(operation,path,"cannot be read",shape)
139	    case True{}: aligned(operation,path,shape,target,count)
140	
141	def loaded(operation: String,path: String,+shape: List<&2,U32>,outcome: Outcome,+count: U32) -> Tensor():
142	  Outcome{target,complete} = outcome
143	  checked(complete,operation,path,shape,target,count)
144	
145	def opened(result: Result<&1,&1,U32 & String,File>,path: String,+shape: List<&2,U32>,+count: U32) -> IO(Tensor()):
146	  match result:
147	    case Fail{error}: IO.pure(Tensor(),fault("load",path,"cannot be opened",shape))
148	    case Done{file}:
149	      do IO<Tensor()>:
150	        outcome : Outcome <- read_all(fuel(count),More{file,fresh(count)},count)
151	        IO.pure(Tensor(),loaded("load",path,shape,outcome,count))
152	
153	def counted(count: Maybe<&2,U32>,+path: String,+shape: List<&2,U32>) -> IO(Tensor()):
154	  match count:
155	    case None{}: IO.pure(Tensor(),Tensors.failed(~F32,"load",Tensors.InvalidShape{shape},shape,0.0,1))
156	    case Some{+count}:
157	      do IO<Tensor()>:
158	        result : Result<&1,&1,U32 & String,File> <- File.open(path,"r")
159	        opened(result,path,shape,count)
160	
161	# The file holds the values of the shape in row-major order and nothing else.
162	def load(+path: String,+shape: List<&2,U32>) -> IO(Tensor()):
163	  counted(Tensors.count_of(shape),path,shape)
164	
165	# The NumPy .npy header: magic, version, header length, then a text such as
166	# {'descr': '<f4', 'fortran_order': False, 'shape': (16, 3, 6, 6), }
167	
168	type Header is Data:
169	  Header{shape: List<&2,U32>,data: List<&2,U32>}
170	
171	def codes(text: String) -> List<&2,U32>:
172	  match text:
173	    case SNil{}: []
174	    case SCon{head,tail}: Char.to_u32(head) <> codes(tail)
175	
176	def magic() -> List<&2,U32>:
177	  [147,78,85,77,80,89]
178	
179	def starts(pattern: List<&2,U32>,bytes: List<&2,U32>) -> Bool:
180	  match pattern bytes:
181	    case Nil{} bytes: True{}
182	    case Con{first,rest} Nil{}: False{}
183	    case Con{first,rest} Con{byte,more}: U32.is_eq(first,byte) && starts(rest,more)
184	
185	def at(bytes: List<&2,U32>,index: Nat) -> U32:
186	  match bytes index:
187	    case Nil{} index: 0
188	    case Con{byte,rest} 0n: byte
189	    case Con{byte,rest} 1n+before: at(rest,before)
190	
191	def contains(text: List<&2,U32>,+pattern: List<&2,U32>) -> Bool:
192	  match text:
193	    case Nil{}: False{}
194	    case Con{+byte,+rest}: starts(pattern,byte <> rest) || contains(rest,pattern)
195	
196	# The text after the first occurrence of the pattern.
197	def after(text: List<&2,U32>,+pattern: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
198	  match text:
199	    case Nil{}: None{}
200	    case Con{+byte,+rest}:
201	      Bool.pick(Maybe<&2,List<&2,U32>>,starts(pattern,byte <> rest),Some{List.drop(&2,U32,byte <> rest,List.length(&2,U32,pattern))},after(rest,pattern))
202	
203	def pushed(digits: Bool,current: U32,shape: List<&2,U32>) -> List<&2,U32>:
204	  match digits:
205	    case True{}: current <> shape
206	    case False{}: shape
207	
208	def closed_shape(valid: Bool,shape: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
209	  match valid:
210	    case True{}: Some{List.reverse(&2,U32,shape)}
211	    case False{}: None{}
212	
213	# Extents up to the closing parenthesis: digits, commas and spaces.
214	def extents(text: List<&2,U32>,+current: U32,digits: Bool,valid: Bool,shape: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
215	  match text:
216	    case Nil{}: None{}
217	    case Con{+byte,rest}:
218	      match byte:
219	        case 41: closed_shape(valid,pushed(digits,current,shape))
220	        case 44: extents(rest,0,False{},valid,pushed(digits,current,shape))
221	        case 32: extents(rest,current,digits,valid,shape)
222	        case _: extents(rest,(current * 10 + byte - 48 : U32),True{},valid && U32.is_le(48,byte) && U32.is_le(byte,57),shape)
223	
224	def shaped(shape: Maybe<&2,List<&2,U32>>,data: List<&2,U32>) -> Maybe<&2,Header>:
225	  match shape:
226	    case None{}: None{}
227	    case Some{shape}: Some{Header{shape,data}}
228	
229	def typed(supported: Bool,text: Maybe<&2,List<&2,U32>>,data: List<&2,U32>) -> Maybe<&2,Header>:
230	  match supported text:
231	    case False{} text: None{}
232	    case True{} None{}: None{}
233	    case True{} Some{text}: shaped(extents(text,0,False{},True{},[]),data)
234	
235	def described(+text: List<&2,U32>,+length: Nat,data: List<&2,U32>) -> Maybe<&2,Header>:
236	  typed(Nat.is_eq(List.length(&2,U32,text),length) && contains(text,codes("'descr': '<f4'")) && contains(text,codes("'fortran_order': False")),
237	    after(text,codes("'shape': (")),data)
238	
239	def region(+bytes: List<&2,U32>,+length: Nat) -> Maybe<&2,Header>:
240	  described(List.take(&2,U32,bytes,length),length,List.drop(&2,U32,bytes,length))
241	
242	# Version 1 has a two-byte header length; later versions have four bytes.
243	def versioned(major: U32,+bytes: List<&2,U32>) -> Maybe<&2,Header>:
244	  match major:
245	    case 1: region(List.drop(&2,U32,bytes,10n),U32.to_nat((at(bytes,8n) + at(bytes,9n) * 256 : U32)))
246	    case _: region(List.drop(&2,U32,bytes,12n),U32.to_nat((at(bytes,8n) + at(bytes,9n) * 256 + at(bytes,10n) * 65536 + at(bytes,11n) * 16777216 : U32)))
247	
248	def marked(valid: Bool,+bytes: List<&2,U32>) -> Maybe<&2,Header>:
249	  match valid:
250	    case False{}: None{}
251	    case True{}: versioned(at(bytes,6n),bytes)
252	
253	def header(+bytes: List<&2,U32>) -> Maybe<&2,Header>:
254	  marked(starts(magic(),bytes),bytes)
255	
256	def refused(file: File,tensor: Tensor()) -> IO(Tensor()):
257	  do IO<Tensor()>:
258	    File.close(file)
259	    IO.pure(Tensor(),tensor)
260	
261	def npy_counted(count: Maybe<&2,U32>,+shape: List<&2,U32>,data: List<&2,U32>,file: File,+path: String) -> IO(Tensor()):
262	  match count:
263	    case None{}: refused(file,Tensors.failed(~F32,"load_npy",Tensors.InvalidShape{shape},shape,0.0,1))
264	    case Some{+count}:
265	      do IO<Tensor()>:
266	        outcome : Outcome <- read_all(fuel(count),More{file,fill(data,fresh(count),count)},count)
267	        IO.pure(Tensor(),loaded("load_npy",path,shape,outcome,count))
268	
269	def npy_parsed(parsed: Maybe<&2,Header>,file: File,+path: String) -> IO(Tensor()):
270	  match parsed:
271	    case None{}: refused(file,fault("load_npy",path,"is not a little-endian FP32 .npy array in C order",[]))
272	    case Some{Header{+shape,data}}: npy_counted(Tensors.count_of(shape),shape,data,file,path)
273	
274	def npy_received(result: Result<&1,&1,U32 & String,List<&2,U32>>,file: File,+path: String) -> IO(Tensor()):
275	  match result:
276	    case Fail{error}: refused(file,fault("load_npy",path,"cannot be read",[]))
277	    case Done{bytes}: npy_parsed(header(bytes),file,path)
278	
279	def npy_headed(read: Read(),+path: String) -> IO(Tensor()):
280	  (file,result) = read
281	  npy_received(result,file,path)
282	
283	def npy_opened(result: Result<&1,&1,U32 & String,File>,+path: String) -> IO(Tensor()):
284	  match result:
285	    case Fail{error}: IO.pure(Tensor(),fault("load_npy",path,"cannot be opened",[]))
286	    case Done{file}:
287	      do IO<Tensor()>:
288	        read : Read() <- File.read_bytes(file,(chunk() * 4 : U32))
289	        npy_headed(read,path)
290	
291	# An FP32 array saved by NumPy, with its shape.
292	def load_npy(+path: String) -> IO(Tensor()):
293	  do IO<Tensor()>:
294	    result : Result<&1,&1,U32 & String,File> <- File.open(path,"r")
295	    npy_opened(result,path)
296	
297	# Writing
298	
299	type Writing is Type:
300	  Pending{file: File,storage: Array<F32>,offset: U32,remaining: Nat}
301	  Finished{file: File,complete: Bool}
302	
303	def bytes(+word: U32,rest: List<&2,U32>) -> List<&2,U32>:
304	  (word % 256 : U32) <> (word / 256 % 256 : U32) <> (word / 65536 % 256 : U32) <> (word / 16777216 : U32) <> rest
305	
306	def encode(count: Nat,values: List<&2,F32>) -> List<&2,U32>:
307	  match count values:
308	    case 0n values: []
309	    case 1n+rest Nil{}: []
310	    case 1n+rest Con{value,more}: bytes(to_bits(value),encode(rest,more))
311	
312	# The cursor retains the array; only the current chunk becomes a list.
313	def read_chunk(count: Nat,+offset: U32,storage: Array<F32>,values: List<&2,F32>) -> Array<F32> & List<&2,F32>:
314	  match count:
315	    case 0n: (storage,List.reverse(&2,F32,values))
316	    case 1n+rest:
317	      Buffers.bind_pair(Array<F32>,F32,Array<F32> & List<&2,F32>,Array.get(F32,storage,offset),storage => value =>
318	        read_chunk(rest,(offset + 1 : U32),storage,value <> values))
319	
320	def pending(remaining: Nat,file: File,storage: Array<F32>,offset: U32) -> Writing:
321	  match remaining:
322	    case 0n: Finished{file,True{}}
323	    case 1n+rest: Pending{file,storage,offset,1n+rest}
324	
325	def continued(result: Result<&1,&1,U32 & String,Unit>,file: File,storage: Array<F32>,remaining: Nat,offset: U32) -> Writing:
326	  match result:
327	    case Fail{error}: Finished{file,False{}}
328	    case Done{unit}: pending(remaining,file,storage,offset)
329	
330	def after_write(written: Written(),storage: Array<F32>,remaining: Nat,offset: U32) -> Writing:
331	  (file,result) = written
332	  continued(result,file,storage,remaining,offset)
333	
334	def write_chunk(file: File,read: Array<F32> & List<&2,F32>,+count: Nat,remaining: Nat,offset: U32) -> IO(Writing):
335	  (storage,values) = read
336	  do IO<Writing>:
337	    written : Written() <- File.write_bytes(file,encode(count,values))
338	    IO.pure(Writing,after_write(written,storage,Nat.sub(remaining,count),(offset + U32.from_nat(count) : U32)))
339	
340	def finished(file: File,complete: Bool) -> IO(Bool):
341	  do IO<Bool>:
342	    File.close(file)
343	    IO.pure(Bool,complete)
344	
345	def write_all(fuel: Nat,writing: Writing) -> IO(Bool):
346	  match fuel writing:
347	    case 0n Pending{file,storage,offset,remaining}: finished(file,False{})
348	    case 0n Finished{file,complete}: finished(file,complete)
349	    case 1n+rest Finished{file,complete}: finished(file,complete)
350	    case 1n+rest Pending{file,storage,+offset,+remaining}:
351	      +count = Nat.min(U32.to_nat(chunk()),remaining)
352	      do IO<Bool>:
353	        next : Writing <- write_chunk(file,read_chunk(count,offset,storage,[]),count,remaining,offset)
354	        write_all(rest,next)
355	
356	def unsaved(pair: Tensor() & List<&2,U32>,operation: String,path: String,detail: String) -> Tensor():
357	  (tensor,shape) = pair
358	  fault(operation,path,detail,shape)
359	
360	def written(complete: Bool,operation: String,path: String,tensor: Tensor()) -> Tensor():
361	  match complete:
362	    case True{}: tensor
363	    case False{}: unsaved(Tensors.shape(~F32,tensor),operation,path,"cannot be written")
364	
365	# The prefix is written first; a raw file has none.
366	def created(result: Result<&1,&1,U32 & String,File>,operation: String,+path: String,tensor: Tensor(),prefix: List<&2,U32>,storage: Array<F32>,+length: Nat) -> IO(Tensor()):
367	  match result:
368	    case Fail{error}: IO.pure(Tensor(),unsaved(Tensors.shape(~F32,tensor),operation,path,"cannot be created"))
369	    case Done{file}:
370	      do IO<Tensor()>:
371	        head : Written() <- File.write_bytes(file,prefix)
372	        complete : Bool <- write_all(Nat.add(length,2n),after_write(head,storage,length,0))
373	        IO.pure(Tensor(),written(complete,operation,path,tensor))
374	
375	def saving(operation: String,+path: String,tensor: Tensor(),prefix: List<&2,U32>,buffer: Result<&2,&1,Tensors.Problem,Buffers.Buffer<F32>>) -> IO(Tensor()):
376	  match buffer:
377	    case Fail{problem}: IO.pure(Tensor(),tensor)
378	    case Done{Buffers.Buffer{storage,length,certificate}}:
379	      do IO<Tensor()>:
380	        result : Result<&1,&1,U32 & String,File> <- File.open(path,"w")
381	        created(result,operation,path,tensor,prefix,storage,length)
382	
383	# Writes the values in row-major order and returns the tensor. A failed tensor
384	# writes nothing and keeps its problem.
385	def save(+path: String,tensor: Tensor()) -> IO(Tensor()):
386	  Tensors.then(~Tensor(),~Tensor(),~IO(Tensor()),Layout.copy(~F32,tensor),kept => copy => saving("save",path,kept,[],Tensors.into_storage(~F32,copy)))
387	
388	def joined(shape: List<&2,U32>) -> String:
389	  match shape:
390	    case []: ""
391	    case [extent]: U32.show(extent)
392	    case extent <> rest: U32.show(extent) ++ ", " ++ joined(rest)
393	
394	# NumPy writes (), (3,) and (2, 3).
395	def tuple(shape: List<&2,U32>) -> String:
396	  match shape:
397	    case []: "()"
398	    case [extent]: "(" ++ U32.show(extent) ++ ",)"
399	    case extent <> rest: "(" ++ joined(extent <> rest) ++ ")"
400	
401	def spaces(count: Nat) -> List<&2,U32>:
402	  match count:
403	    case 0n: [10]
404	    case 1n+rest: 32 <> spaces(rest)
405	
406	# Version 1: the header ends in a newline and the data starts at a multiple of 64 bytes.
407	def framed(+text: List<&2,U32>,+length: U32) -> List<&2,U32>:
408	  +padding = ((64 - (11 + length) % 64) % 64 : U32)
409	  +total = (length + padding + 1 : U32)
410	  List.append(&2,U32,magic(),1 <> 0 <> (total % 256 : U32) <> (total / 256 : U32) <> List.append(&2,U32,text,spaces(U32.to_nat(padding))))
411	
412	def npy_header(shape: List<&2,U32>) -> List<&2,U32>:
413	  +text = codes("{'descr': '<f4', 'fortran_order': False, 'shape': " ++ tuple(shape) ++ ", }")
414	  framed(text,U32.from_nat(List.length(&2,U32,text)))
415	
416	def npy_saving(+path: String,pair: Tensor() & Tensor(),+shape: List<&2,U32>) -> IO(Tensor()):
417	  (kept,copy) = pair
418	  saving("save_npy",path,kept,npy_header(shape),Tensors.into_storage(~F32,copy))
419	
420	# Writes a .npy file that NumPy loads with the tensor's shape.
421	def save_npy(+path: String,tensor: Tensor()) -> IO(Tensor()):
422	  Tensors.then(~Tensor(),~List<&2,U32>,~IO(Tensor()),Tensors.shape(~F32,tensor),tensor => shape => npy_saving(path,Layout.copy(~F32,tensor),shape))