1	# Exact overflow-checked multiplication for metadata. Each Horner step starts
2	# below word_max and forms at most 3*word_max, safely below the native Nat cap.
3	# The loop has exactly 32 bit steps and never touches tensor payloads.
4	import Base
5	import ./machine_limits.bend as Machine
6	
7	def bounded_result(valid: Bool,value: Nat) -> Maybe<&2,Nat>:
8	  match valid:
9	    case False{}: None{}
10	    case True{}: Some{value}
11	
12	def bounded(+value: Nat,limit: Nat) -> Maybe<&2,Nat>:
13	  bounded_result(Nat.is_le(value,limit),value)
14	
15	def multiply_step(bit: Bool,right: Nat,previous: Maybe<&2,Nat>,limit: Nat) -> Maybe<&2,Nat>:
16	  match previous:
17	    case None{}: None{}
18	    case Some{value}: bounded(Nat.add(Nat.double(value),Bool.pick(Nat,bit,right,0n)),limit)
19	
20	def multiply_bits(width: Nat,bits: Word(width),+right: Nat,+limit: Nat) -> Maybe<&2,Nat>:
21	  match width bits:
22	    case 0n WNil{}: Some{0n}
23	    case 1n+rest WCon{bit,tail}: multiply_step(bit,right,multiply_bits(rest,tail,right,limit),limit)
24	
25	def encoded(result: Maybe<&2,Nat>) -> Maybe<&2,U32>:
26	  match result:
27	    case None{}: None{}
28	    case Some{value}: Some{U32.from_nat(value)}
29	
30	def multiply_bounded(left: U32,right: U32,limit: U32) -> Maybe<&2,U32>:
31	  U32{bits} = left
32	  encoded(multiply_bits(32n,bits,U32.to_nat(right),U32.to_nat(limit)))
33	
34	# A left operand at most 2^31 divided by a right operand of at most 2^30
35	# gives a product of at most 2^31, which the native product holds. Other
36	# products take the exact bit traversal.
37	def multiply_small(fits: Bool,+left: U32,+right: U32) -> Maybe<&2,U32>:
38	  match fits:
39	    case True{}: Some{(left * right : U32)}
40	    case False{}: multiply_bounded(left,right,Machine.word_max())
41	
42	def multiply_divided(large: Bool,+left: U32,+right: U32) -> Maybe<&2,U32>:
43	  match large:
44	    case True{}: multiply_bounded(left,right,Machine.word_max())
45	    case False{}: multiply_small(U32.is_le(left,(2147483648 / right : U32)),left,right)
46	
47	# Unit products already fit a word; shapes and contiguous strides frequently
48	# use them.
49	def multiply_right(unit: Bool,+left: U32,+right: U32) -> Maybe<&2,U32>:
50	  match unit:
51	    case True{}: Some{left}
52	    case False{}: multiply_divided(U32.is_zero(right) || U32.is_lt(1073741824,right),left,right)
53	
54	def multiply_left(unit: Bool,+left: U32,+right: U32) -> Maybe<&2,U32>:
55	  match unit:
56	    case True{}: Some{right}
57	    case False{}: multiply_right(U32.is_eq(right,1),left,right)
58	
59	def multiply(+left: U32,+right: U32) -> Maybe<&2,U32>:
60	  multiply_left(U32.is_eq(left,1),left,right)