1	# Storage-independent extents. View derives this descriptor from its axes.
2	import Base
3	import ./math_checked_word.bend as CheckedWord
4	
5	type Shape is Data:
6	  Shape{extents: List<&2,U32>}
7	
8	def rank(shape: Shape) -> Nat:
9	  Shape{extents} = shape
10	  List.length(&2,U32,extents)
11	
12	def extent(shape: Shape,axis: Nat) -> Maybe<&2,U32>:
13	  Shape{extents} = shape
14	  List.get(&2,U32,extents,axis)
15	
16	def same_extents(left: List<&2,U32>,right: List<&2,U32>) -> Bool:
17	  match left right:
18	    case Nil{} Nil{}: True{}
19	    case Con{x,xs} Con{y,ys}: U32.is_eq(x,y) && same_extents(xs,ys)
20	    case left right: False{}
21	
22	def same(left: Shape,right: Shape) -> Bool:
23	  Shape{xs} = left
24	  Shape{ys} = right
25	  same_extents(xs,ys)
26	
27	def has_zero(extents: List<&2,U32>) -> Bool:
28	  match extents:
29	    case Nil{}: False{}
30	    case Con{x,xs}: U32.is_zero(x) || has_zero(xs)
31	
32	def checked_multiply(left: U32,right: U32) -> Maybe<&2,U32>:
33	  CheckedWord.multiply(left,right)
34	
35	def product(extents: List<&2,U32>,accumulator: Maybe<&2,U32>) -> Maybe<&2,U32>:
36	  match extents accumulator:
37	    case Nil{} result: result
38	    case Con{x,xs} Some{value}: product(xs,checked_multiply(value,x))
39	    case Con{x,xs} None{}: None{}
40	
41	def count_nonempty(empty: Bool,extents: List<&2,U32>) -> Maybe<&2,U32>:
42	  match empty:
43	    case True{}: Some{0}
44	    case False{}: product(extents,Some{1})
45	
46	# Scalar [] has one element. Any zero extent makes an empty shape, even when
47	# an earlier partial product would overflow. Nonempty overflow is rejected.
48	def numel(shape: Shape) -> Maybe<&2,U32>:
49	  Shape{+extents} = shape
50	  count_nonempty(has_zero(extents),extents)
51	
52	# NumPy broadcasting of one axis: equal extents, or one of them is 1.
53	def broadcast_extent(same: Bool,left_one: Bool,right_one: Bool,left: U32,right: U32) -> Maybe<&2,U32>:
54	  match same left_one right_one:
55	    case True{} left_one right_one: Some{left}
56	    case False{} True{} right_one: Some{right}
57	    case False{} False{} True{}: Some{left}
58	    case False{} False{} False{}: None{}
59	
60	def prepend_extent(extent: Maybe<&2,U32>,rest: Maybe<&2,List<&2,U32>>) -> Maybe<&2,List<&2,U32>>:
61	  match extent rest:
62	    case Some{extent} Some{rest}: Some{extent <> rest}
63	    case extent rest: None{}
64	
65	def broadcast_same_rank(left: List<&2,U32>,right: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
66	  match left right:
67	    case Nil{} Nil{}: Some{[]}
68	    case Con{+x,xs} Con{+y,ys}:
69	      prepend_extent(broadcast_extent(U32.is_eq(x,y),U32.is_eq(x,1),U32.is_eq(y,1),x,y),broadcast_same_rank(xs,ys))
70	    case left right: None{}
71	
72	# The extents with leading axes of extent 1 added up to rank.
73	def with_rank(+extents: List<&2,U32>,rank: Nat) -> List<&2,U32>:
74	  List.append(&2,U32,List.replicate(U32,Nat.sub(rank,List.length(&2,U32,extents)),1),extents)
75	
76	# The common extents of two operands under NumPy broadcasting: the extents are
77	# aligned at the last axis, a missing axis counts as 1, and on each axis the
78	# extents are equal or one of them is 1.
79	def broadcast(+left: List<&2,U32>,+right: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
80	  broadcast_same_rank(with_rank(left,List.length(&2,U32,right)),with_rank(right,List.length(&2,U32,left)))
81	
82	def sum_when(valid: Bool,value: U32) -> Maybe<&2,U32>:
83	  match valid:
84	    case True{}: Some{value}
85	    case False{}: None{}
86	
87	def checked_add(+left: U32,+right: U32) -> Maybe<&2,U32>:
88	  sum_when(U32.is_le(right,(4294967295 - left : U32)),(left + right : U32))
89	
90	def checked_add_to(left: Maybe<&2,U32>,right: U32) -> Maybe<&2,U32>:
91	  match left:
92	    case Some{left}: checked_add(left,right)
93	    case None{}: None{}
94	
95	# Each extent with before elements added in front and after elements behind.
96	# The three lists have one entry per axis and no sum may wrap.
97	def padded(extents: List<&2,U32>,before: List<&2,U32>,after: List<&2,U32>) -> Maybe<&2,List<&2,U32>>:
98	  match extents before after:
99	    case Nil{} Nil{} Nil{}: Some{[]}
100	    case Con{extent,extents} Con{front,before} Con{back,after}:
101	      prepend_extent(checked_add_to(checked_add(front,extent),back),padded(extents,before,after))
102	    case extents before after: None{}