1use crate::arithmetic::{fp_div_i, fp_mul_i, fp_sqrt};
17use crate::bvn_cdf::bvn_cdf;
18use crate::constants::{SCALE, SCALE_I};
19use crate::error::SolMathError;
20use crate::transcendental::{exp_fixed_i, ln_fixed_i};
21
22const MAX_INPUT: u128 = 100_000 * SCALE;
23
24struct RainbowTerms {
26 s1: i128,
27 s2: i128,
28 disc_q1: i128, disc_q2: i128, disc_r_k: i128, y1: i128,
32 y2: i128,
33 d: i128,
34 srt: i128, v1_sqrt_t: i128,
36 v2_sqrt_t: i128,
37 r1: i128, r2: i128, rho: i128,
40}
41
42#[allow(clippy::too_many_arguments)]
43fn prepare(
44 s1: u128,
45 s2: u128,
46 k: u128,
47 r: u128,
48 q1: u128,
49 q2: u128,
50 sigma1: u128,
51 sigma2: u128,
52 rho: i128,
53 t: u128,
54) -> Result<RainbowTerms, SolMathError> {
55 if s1 > MAX_INPUT
56 || s2 > MAX_INPUT
57 || k > MAX_INPUT
58 || r > MAX_INPUT
59 || q1 > MAX_INPUT
60 || q2 > MAX_INPUT
61 || sigma1 > MAX_INPUT
62 || sigma2 > MAX_INPUT
63 || t > MAX_INPUT
64 {
65 return Err(SolMathError::Overflow);
66 }
67 if s1 == 0 || s2 == 0 || k == 0 || sigma1 == 0 || sigma2 == 0 || t == 0 {
68 return Err(SolMathError::DomainError);
69 }
70 if rho.unsigned_abs() >= SCALE {
71 return Err(SolMathError::DomainError);
72 }
73 let (s1, s2, k, r) = (s1 as i128, s2 as i128, k as i128, r as i128);
74 let (q1, q2, sigma1, sigma2, t) = (
75 q1 as i128,
76 q2 as i128,
77 sigma1 as i128,
78 sigma2 as i128,
79 t as i128,
80 );
81 let b1 = r - q1;
82 let b2 = r - q2;
83
84 let sqrt_t = fp_sqrt(t as u128)? as i128;
85 let s1_sq = fp_mul_i(sigma1, sigma1)?;
86 let s2_sq = fp_mul_i(sigma2, sigma2)?;
87 let cross = fp_mul_i(fp_mul_i(rho, sigma1)?, sigma2)?;
89 let sig_sq = s1_sq + s2_sq - 2 * cross;
90 if sig_sq <= 0 {
91 return Err(SolMathError::DomainError);
92 }
93 let sig = fp_sqrt(sig_sq as u128)? as i128;
94 let srt = fp_mul_i(sig, sqrt_t)?;
95 let v1_sqrt_t = fp_mul_i(sigma1, sqrt_t)?;
96 let v2_sqrt_t = fp_mul_i(sigma2, sqrt_t)?;
97 if srt == 0 || v1_sqrt_t == 0 || v2_sqrt_t == 0 {
98 return Err(SolMathError::DomainError);
99 }
100
101 let ln_s1s2 = ln_fixed_i(fp_div_i(s1, s2)? as u128)?;
103 let d = fp_div_i(ln_s1s2 + fp_mul_i(b1 - b2 + sig_sq / 2, t)?, srt)?;
104 let y1 = fp_div_i(
106 ln_fixed_i(fp_div_i(s1, k)? as u128)? + fp_mul_i(b1 + s1_sq / 2, t)?,
107 v1_sqrt_t,
108 )?;
109 let y2 = fp_div_i(
110 ln_fixed_i(fp_div_i(s2, k)? as u128)? + fp_mul_i(b2 + s2_sq / 2, t)?,
111 v2_sqrt_t,
112 )?;
113 let r1 = fp_div_i(fp_mul_i(rho, sigma2)? - sigma1, sig)?;
115 let r2 = fp_div_i(fp_mul_i(rho, sigma1)? - sigma2, sig)?;
116
117 Ok(RainbowTerms {
118 s1,
119 s2,
120 disc_q1: exp_fixed_i(-fp_mul_i(q1, t)?)?,
121 disc_q2: exp_fixed_i(-fp_mul_i(q2, t)?)?,
122 disc_r_k: fp_mul_i(k, exp_fixed_i(-fp_mul_i(r, t)?)?)?,
123 y1,
124 y2,
125 d,
126 srt,
127 v1_sqrt_t,
128 v2_sqrt_t,
129 r1,
130 r2,
131 rho,
132 })
133}
134
135#[allow(clippy::too_many_arguments)]
140pub fn worst_of_call(
141 s1: u128,
142 s2: u128,
143 k: u128,
144 r: u128,
145 q1: u128,
146 q2: u128,
147 sigma1: u128,
148 sigma2: u128,
149 rho: i128,
150 t: u128,
151) -> Result<u128, SolMathError> {
152 let tm = prepare(s1, s2, k, r, q1, q2, sigma1, sigma2, rho, t)?;
153 let a = fp_mul_i(fp_mul_i(tm.s1, tm.disc_q1)?, bvn_cdf(tm.y1, -tm.d, tm.r1)?)?;
156 let b = fp_mul_i(
157 fp_mul_i(tm.s2, tm.disc_q2)?,
158 bvn_cdf(tm.y2, tm.d - tm.srt, tm.r2)?,
159 )?;
160 let c = fp_mul_i(
161 tm.disc_r_k,
162 bvn_cdf(tm.y1 - tm.v1_sqrt_t, tm.y2 - tm.v2_sqrt_t, tm.rho)?,
163 )?;
164 Ok((a + b - c).max(0) as u128)
165}
166
167#[allow(clippy::too_many_arguments)]
171pub fn best_of_call(
172 s1: u128,
173 s2: u128,
174 k: u128,
175 r: u128,
176 q1: u128,
177 q2: u128,
178 sigma1: u128,
179 sigma2: u128,
180 rho: i128,
181 t: u128,
182) -> Result<u128, SolMathError> {
183 let tm = prepare(s1, s2, k, r, q1, q2, sigma1, sigma2, rho, t)?;
184 let a = fp_mul_i(fp_mul_i(tm.s1, tm.disc_q1)?, bvn_cdf(tm.y1, tm.d, -tm.r1)?)?;
187 let b = fp_mul_i(
188 fp_mul_i(tm.s2, tm.disc_q2)?,
189 bvn_cdf(tm.y2, tm.srt - tm.d, -tm.r2)?,
190 )?;
191 let m = bvn_cdf(-tm.y1 + tm.v1_sqrt_t, -tm.y2 + tm.v2_sqrt_t, tm.rho)?;
192 let c = fp_mul_i(tm.disc_r_k, SCALE_I - m)?;
193 Ok((a + b - c).max(0) as u128)
194}