Skip to content

Commit ffa533f

Browse files
jordy25519jordy25519
committed
wip
1 parent 2479d69 commit ffa533f

File tree

5 files changed

+272
-0
lines changed

5 files changed

+272
-0
lines changed

‎programs/drift/src/math/bn.rs‎

Lines changed: 254 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ use crate::error::DriftResult;
1414

1515
pub mod compat {
1616
use bytemuck::{Pod, Zeroable};
17+
use num_traits::{One, Zero};
18+
use std::cmp::{Ord, Ordering, PartialOrd};
1719
use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign};
1820

1921
/// `u128` with legacy bit layout
@@ -114,6 +116,131 @@ pub mod compat {
114116
}
115117
}
116118

119+
// Comparison traits for u128
120+
impl PartialOrd for u128 {
121+
#[inline(always)]
122+
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
123+
self.as_u128().partial_cmp(&other.as_u128())
124+
}
125+
}
126+
127+
impl Ord for u128 {
128+
#[inline(always)]
129+
fn cmp(&self, other: &Self) -> Ordering {
130+
self.as_u128().cmp(&other.as_u128())
131+
}
132+
}
133+
134+
// Helper methods for u128
135+
impl u128 {
136+
/// Returns the absolute value (no-op for unsigned)
137+
#[inline(always)]
138+
pub fn abs(self) -> Self {
139+
self
140+
}
141+
142+
/// Returns the signum (1 if > 0, 0 if == 0)
143+
#[inline(always)]
144+
pub fn signum(self) -> Self {
145+
if self.as_u128() > 0 {
146+
Self::ONE
147+
} else {
148+
Self::ZERO
149+
}
150+
}
151+
152+
/// Returns the minimum of two values
153+
#[inline(always)]
154+
pub fn min(self, other: Self) -> Self {
155+
if self < other {
156+
self
157+
} else {
158+
other
159+
}
160+
}
161+
162+
/// Returns the maximum of two values
163+
#[inline(always)]
164+
pub fn max(self, other: Self) -> Self {
165+
if self > other {
166+
self
167+
} else {
168+
other
169+
}
170+
}
171+
172+
/// Saturating addition
173+
#[inline(always)]
174+
pub fn saturating_add(self, other: Self) -> Self {
175+
Self::from(self.as_u128().saturating_add(other.as_u128()))
176+
}
177+
178+
/// Saturating subtraction
179+
#[inline(always)]
180+
pub fn saturating_sub(self, other: Self) -> Self {
181+
Self::from(self.as_u128().saturating_sub(other.as_u128()))
182+
}
183+
184+
/// Saturating multiplication
185+
#[inline(always)]
186+
pub fn saturating_mul(self, other: Self) -> Self {
187+
Self::from(self.as_u128().saturating_mul(other.as_u128()))
188+
}
189+
190+
/// Checked addition
191+
#[inline(always)]
192+
pub fn checked_add(self, other: Self) -> Option<Self> {
193+
self.as_u128().checked_add(other.as_u128()).map(Self::from)
194+
}
195+
196+
/// Checked subtraction
197+
#[inline(always)]
198+
pub fn checked_sub(self, other: Self) -> Option<Self> {
199+
self.as_u128().checked_sub(other.as_u128()).map(Self::from)
200+
}
201+
202+
/// Checked multiplication
203+
#[inline(always)]
204+
pub fn checked_mul(self, other: Self) -> Option<Self> {
205+
self.as_u128().checked_mul(other.as_u128()).map(Self::from)
206+
}
207+
208+
/// Checked division
209+
#[inline(always)]
210+
pub fn checked_div(self, other: Self) -> Option<Self> {
211+
self.as_u128().checked_div(other.as_u128()).map(Self::from)
212+
}
213+
214+
/// Checked remainder
215+
#[inline(always)]
216+
pub fn checked_rem(self, other: Self) -> Option<Self> {
217+
self.as_u128().checked_rem(other.as_u128()).map(Self::from)
218+
}
219+
}
220+
221+
// num_traits implementations for u128
222+
impl Zero for u128 {
223+
fn zero() -> Self {
224+
Self::ZERO
225+
}
226+
227+
fn is_zero(&self) -> bool {
228+
*self == Self::ZERO
229+
}
230+
}
231+
232+
impl One for u128 {
233+
fn one() -> Self {
234+
Self::ONE
235+
}
236+
}
237+
238+
impl std::fmt::Display for self::u128 {
239+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
240+
self.as_u128().fmt(f)
241+
}
242+
}
243+
117244
/// `i128` with legacy bit layout
118245
#[derive(Copy, Clone, PartialEq, Eq, Debug, Default)]
119246
#[repr(transparent)]
@@ -215,6 +342,133 @@ pub mod compat {
215342
*self = *self / other;
216343
}
217344
}
345+
346+
// Comparison traits for i128
347+
impl PartialOrd for i128 {
348+
#[inline(always)]
349+
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
350+
self.as_i128().partial_cmp(&other.as_i128())
351+
}
352+
}
353+
354+
impl Ord for i128 {
355+
#[inline(always)]
356+
fn cmp(&self, other: &Self) -> Ordering {
357+
self.as_i128().cmp(&other.as_i128())
358+
}
359+
}
360+
361+
// Helper methods for i128
362+
impl i128 {
363+
/// Returns the absolute value
364+
#[inline(always)]
365+
pub fn abs(self) -> Self {
366+
Self::from(self.as_i128().abs())
367+
}
368+
369+
/// Returns the signum (-1, 0, or 1)
370+
#[inline(always)]
371+
pub fn signum(self) -> Self {
372+
Self::from(self.as_i128().signum())
373+
}
374+
375+
/// Returns the absolute value as unsigned
376+
#[inline(always)]
377+
pub fn unsigned_abs(self) -> u128 {
378+
u128::from(self.as_i128().unsigned_abs())
379+
}
380+
381+
/// Returns the minimum of two values
382+
#[inline(always)]
383+
pub fn min(self, other: Self) -> Self {
384+
if self < other {
385+
self
386+
} else {
387+
other
388+
}
389+
}
390+
391+
/// Returns the maximum of two values
392+
#[inline(always)]
393+
pub fn max(self, other: Self) -> Self {
394+
if self > other {
395+
self
396+
} else {
397+
other
398+
}
399+
}
400+
401+
/// Saturating addition
402+
#[inline(always)]
403+
pub fn saturating_add(self, other: Self) -> Self {
404+
Self::from(self.as_i128().saturating_add(other.as_i128()))
405+
}
406+
407+
/// Saturating subtraction
408+
#[inline(always)]
409+
pub fn saturating_sub(self, other: Self) -> Self {
410+
Self::from(self.as_i128().saturating_sub(other.as_i128()))
411+
}
412+
413+
/// Saturating multiplication
414+
#[inline(always)]
415+
pub fn saturating_mul(self, other: Self) -> Self {
416+
Self::from(self.as_i128().saturating_mul(other.as_i128()))
417+
}
418+
419+
/// Checked addition
420+
#[inline(always)]
421+
pub fn checked_add(self, other: Self) -> Option<Self> {
422+
self.as_i128().checked_add(other.as_i128()).map(Self::from)
423+
}
424+
425+
/// Checked subtraction
426+
#[inline(always)]
427+
pub fn checked_sub(self, other: Self) -> Option<Self> {
428+
self.as_i128().checked_sub(other.as_i128()).map(Self::from)
429+
}
430+
431+
/// Checked multiplication
432+
#[inline(always)]
433+
pub fn checked_mul(self, other: Self) -> Option<Self> {
434+
self.as_i128().checked_mul(other.as_i128()).map(Self::from)
435+
}
436+
437+
/// Checked division
438+
#[inline(always)]
439+
pub fn checked_div(self, other: Self) -> Option<Self> {
440+
self.as_i128().checked_div(other.as_i128()).map(Self::from)
441+
}
442+
443+
/// Checked remainder
444+
#[inline(always)]
445+
pub fn checked_rem(self, other: Self) -> Option<Self> {
446+
self.as_i128().checked_rem(other.as_i128()).map(Self::from)
447+
}
448+
}
449+
450+
// num_traits implementations for i128
451+
impl Zero for i128 {
452+
fn zero() -> Self {
453+
Self::ZERO
454+
}
455+
456+
fn is_zero(&self) -> bool {
457+
*self == Self::ZERO
458+
}
459+
}
460+
461+
impl One for i128 {
462+
fn one() -> Self {
463+
Self::ONE
464+
}
465+
}
466+
467+
impl std::fmt::Display for self::i128 {
468+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
469+
self.as_i128().fmt(f)
470+
}
471+
}
218472
}
219473

220474
construct_uint! {

‎programs/drift/src/math/casting.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,3 +36,7 @@ impl Cast for i32 {}
3636
impl Cast for i16 {}
3737
impl Cast for i8 {}
3838
impl Cast for bool {}
39+
40+
// Cast trait for custom compat types
41+
impl Cast for crate::math::bn::compat::u128 {}
42+
impl Cast for crate::math::bn::compat::i128 {}

‎programs/drift/src/math/ceil_div.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,3 +38,7 @@ checked_impl!(i64);
3838
checked_impl!(i32);
3939
checked_impl!(i16);
4040
checked_impl!(i8);
41+
42+
// CheckedCeilDiv for custom compat types
43+
checked_impl!(crate::math::bn::compat::u128);
44+
checked_impl!(crate::math::bn::compat::i128);

‎programs/drift/src/math/floor_div.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@ checked_impl!(i32);
3131
checked_impl!(i16);
3232
checked_impl!(i8);
3333

34+
// CheckedFloorDiv for custom compat i128
35+
checked_impl!(crate::math::bn::compat::i128);
36+
3437
#[cfg(test)]
3538
mod test {
3639
use crate::math::floor_div::CheckedFloorDiv;

‎programs/drift/src/math/safe_math.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,10 @@ checked_impl!(i32);
9797
checked_impl!(i16);
9898
checked_impl!(i8);
9999

100+
// SafeMath for custom compat types
101+
checked_impl!(crate::math::bn::compat::u128);
102+
checked_impl!(crate::math::bn::compat::i128);
103+
100104
pub trait SafeDivFloor: Sized {
101105
/// Perform floor division
102106
fn safe_div_floor(self, rhs: Self) -> DriftResult<Self>;
@@ -127,6 +131,9 @@ div_floor_impl!(i32);
127131
div_floor_impl!(i16);
128132
div_floor_impl!(i8);
129133

134+
// SafeDivFloor for custom compat i128
135+
div_floor_impl!(crate::math::bn::compat::i128);
136+
130137
#[cfg(test)]
131138
mod test {
132139
use crate::error::ErrorCode;

0 commit comments

Comments
 (0)