Add saturating integer math (#7)

Author: BioTurboNick

src/OverflowContexts.jl

@@ -2,10 +2,12 @@ module OverflowContexts
 
 include("macros.jl")
 include("base_ext.jl")
+include("base_ext_sat.jl")
 include("abstractarraymath_ext.jl")
 
-export @default_checked, @default_unchecked, @checked, @unchecked,
+export @default_checked, @default_unchecked, @default_saturating, @checked, @unchecked, @saturating,
+    checked_neg, checked_add, checked_sub, checked_mul, checked_pow, checked_negsub, checked_abs,
     unchecked_neg, unchecked_add, unchecked_sub, unchecked_mul, unchecked_negsub, unchecked_pow, unchecked_abs,
-    checked_neg, checked_add, checked_sub, checked_mul, checked_pow, checked_negsub, checked_abs
+    saturating_neg, saturating_add, saturating_sub, saturating_mul, saturating_pow, saturating_negsub, saturating_abs
 
 end # module

src/base_ext.jl

@@ -1,11 +1,13 @@
-using Base: promote, afoldl, @_inline_meta
+using Base: BitInteger, promote, afoldl, @_inline_meta
 import Base.Checked: checked_neg, checked_add, checked_sub, checked_mul, checked_abs  
+using Base.Checked: mul_with_overflow
 
 if VERSION ≥ v"1.11-alpha"
+    import Base: power_by_squaring
     import Base.Checked: checked_pow
 else
-    using Base: BitInteger, throw_domerr_powbysq, to_power_type
-    using Base.Checked: mul_with_overflow, throw_overflowerr_binaryop
+    using Base: throw_domerr_powbysq, to_power_type
+    using Base.Checked: throw_overflowerr_binaryop
 end
 
 # The Base methods have unchecked semantics, so just pass through
@@ -22,13 +24,22 @@ checked_add(a, b, c, xs...) = @checked (@_inline_meta; afoldl(+, (+)((+)(a, b),
 checked_sub(a, b, c, xs...) = @checked (@_inline_meta; afoldl(-, (-)((-)(a, b), c), xs...))
 checked_mul(a, b, c, xs...) = @checked (@_inline_meta; afoldl(*, (*)((*)(a, b), c), xs...))
 
+saturating_add(a, b, c, xs...) = @saturating (@_inline_meta; afoldl(+, (+)((+)(a, b), c), xs...))
+saturating_sub(a, b, c, xs...) = @saturating (@_inline_meta; afoldl(-, (-)((-)(a, b), c), xs...))
+saturating_mul(a, b, c, xs...) = @saturating (@_inline_meta; afoldl(*, (*)((*)(a, b), c), xs...))
+
 
 # promote unmatched number types to same type
 checked_add(x::Number, y::Number) = checked_add(promote(x, y)...)
 checked_sub(x::Number, y::Number) = checked_sub(promote(x, y)...)
 checked_mul(x::Number, y::Number) = checked_mul(promote(x, y)...)
 checked_pow(x::Number, y::Number) = checked_pow(promote(x, y)...)
 
+saturating_add(x::Number, y::Number) = saturating_add(promote(x, y)...)
+saturating_sub(x::Number, y::Number) = saturating_sub(promote(x, y)...)
+saturating_mul(x::Number, y::Number) = saturating_mul(promote(x, y)...)
+saturating_pow(x::Number, y::Number) = saturating_pow(promote(x, y)...)
+
 
 # fallback to `unchecked_` for `Number` types that don't have more specific `checked_` methods
 checked_neg(x::T) where T <: Number = unchecked_neg(x)
@@ -38,6 +49,13 @@ checked_mul(x::T, y::T) where T <: Number = unchecked_mul(x, y)
 checked_pow(x::T, y::T) where T <: Number = unchecked_pow(x, y)
 checked_abs(x::T) where T <: Number = unchecked_abs(x)
 
+saturating_neg(x::T) where T <: Number = unchecked_neg(x)
+saturating_add(x::T, y::T) where T <: Number = unchecked_add(x, y)
+saturating_sub(x::T, y::T) where T <: Number = unchecked_sub(x, y)
+saturating_mul(x::T, y::T) where T <: Number = unchecked_mul(x, y)
+saturating_pow(x::T, y::T) where T <: Number = unchecked_pow(x, y)
+saturating_abs(x::T) where T <: Number = unchecked_abs(x)
+
 
 # fallback to `unchecked_` for non-`Number` types
 checked_neg(x) = unchecked_neg(x)
@@ -51,50 +69,38 @@ checked_abs(x) = unchecked_abs(x)
 if VERSION < v"1.11"
 # Base.Checked only gained checked powers in 1.11
 
-function checked_pow(x::T, y::S) where {T <: BitInteger, S <: BitInteger}
-    @_inline_meta
-    z, b = pow_with_overflow(x, y)
-    b && throw_overflowerr_binaryop(:^, x, y)
-    z
-end
+checked_pow(x_::T, p::S) where {T <: BitInteger, S <: BitInteger} =
+    power_by_squaring(x_, p; mul = checked_mul)
 
-function pow_with_overflow(x_, p::Integer)
+# Base.@assume_effects :terminates_locally # present in Julia 1.11 code, but only supported from 1.8 on
+function power_by_squaring(x_, p::Integer; mul=*)
     x = to_power_type(x_)
     if p == 1
-        return (copy(x), false)
+        return copy(x)
     elseif p == 0
-        return (one(x), false)
+        return one(x)
     elseif p == 2
-        return mul_with_overflow(x, x)
+        return mul(x, x)
     elseif p < 0
-        isone(x) && return (copy(x), false)
-        isone(-x) && return (iseven(p) ? one(x) : copy(x), false)
+        isone(x) && return copy(x)
+        isone(-x) && return iseven(p) ? one(x) : copy(x)
         throw_domerr_powbysq(x, p)
     end
     t = trailing_zeros(p) + 1
     p >>= t
-    b = false
     while (t -= 1) > 0
-        x, b1 = mul_with_overflow(x, x)
-        b |= b1
+        x = mul(x, x)
     end
     y = x
     while p > 0
         t = trailing_zeros(p) + 1
         p >>= t
         while (t -= 1) >= 0
-            x, b1 = mul_with_overflow(x, x)
-            b |= b1
+            x = mul(x, x)
         end
-        y, b1 = mul_with_overflow(y, x)
-        b |= b1
+        y = mul(y, x)
     end
-    return y, b
-end
-pow_with_overflow(x::Bool, p::Unsigned) = ((p==0) | x, false)
-function pow_with_overflow(x::Bool, p::Integer)
-    p < 0 && !x && throw_domerr_powbysq(x, p)
-    return (p==0) | x, false
+    return y
 end
 
 end

src/base_ext_sat.jl

@@ -0,0 +1,70 @@
+import Base: BitInteger
+import Base.Checked: mul_with_overflow
+
+if VERSION ≤ v"1.11-alpha"
+    import Base: power_by_squaring
+end
+
+# saturating implementations
+const SignedBitInteger = Union{Int8, Int16, Int32, Int64, Int128}
+
+saturating_neg(x::T) where T <: BitInteger = saturating_sub(zero(T), x)
+
+if VERSION ≥ v"1.5"
+    using Base: llvmcall
+
+    # These intrinsics were added in LLVM 8, which was first supported with Julia 1.5
+    @generated function saturating_add(x::T, y::T) where T <: BitInteger
+        llvm_su = T <: Signed ? "s" : "u"
+        llvm_t = "i" * string(8sizeof(T))
+        llvm_intrinsic = "llvm.$(llvm_su)add.sat.$llvm_t"
+        :(ccall($llvm_intrinsic, llvmcall, $T, ($T, $T), x, y))
+    end
+
+    @generated function saturating_sub(x::T, y::T) where T <: BitInteger
+        llvm_su = T <: Signed ? "s" : "u"
+        llvm_t = "i" * string(8sizeof(T))
+        llvm_intrinsic = "llvm.$(llvm_su)sub.sat.$llvm_t"
+        :(ccall($llvm_intrinsic, llvmcall, $T, ($T, $T), x, y))
+    end
+
+else
+    import Base.Checked: add_with_overflow, sub_with_overflow
+
+    function saturating_add(x::T, y::T) where T <: BitInteger
+        result, overflow_flag = add_with_overflow(x, y)
+        if overflow_flag
+            return sign(x) > 0 ?
+                typemax(T) :
+                typemin(T)
+        end
+        return result
+    end
+
+    function saturating_sub(x::T, y::T) where T <: BitInteger
+        result, overflow_flag = sub_with_overflow(x, y)
+        if overflow_flag
+            return y > x ?
+                typemin(T) :
+                typemax(T)
+        end
+        return result
+    end
+end
+
+function saturating_mul(x::T, y::T) where T <: BitInteger
+    result, overflow_flag = mul_with_overflow(x, y)
+    return overflow_flag ?
+        (sign(x) == sign(y) ?
+            typemax(T) :
+            typemin(T)) :
+        result
+end
+
+saturating_pow(x_::T, p::S) where {T <: BitInteger, S <: BitInteger} =
+    power_by_squaring(x_, p; mul = saturating_mul)
+
+function saturating_abs(x::T) where T <: SignedBitInteger 
+    result = flipsign(x, x)
+    return result < 0 ? typemax(T) : result
+end

src/macros.jl

@@ -50,6 +50,41 @@ macro default_unchecked()
     end
 end
 
+"""
+    @default_saturating
+
+Redirect default integer math to saturating operators for the current module. Only works at top-level.
+"""
+macro default_saturating()
+    quote
+        if !isdefined(@__MODULE__, :__OverflowContextDefaultSet)
+            any(Base.isbindingresolved.(Ref(@__MODULE__), op_method_symbols)) &&
+                error("A default context may only be set before any reference to the affected methods (+, -, *, ^, abs) in the target module.")
+        else
+            @warn "A previous default was set for this module. Previously defined methods in this module will be recompiled with this new default."
+        end
+        (@__MODULE__).eval(:(-(x) = OverflowContexts.saturating_neg(x)))
+        (@__MODULE__).eval(:(+(x...) = OverflowContexts.saturating_add(x...)))
+        (@__MODULE__).eval(:(-(x...) = OverflowContexts.saturating_sub(x...)))
+        (@__MODULE__).eval(:(*(x...) = OverflowContexts.saturating_mul(x...)))
+        (@__MODULE__).eval(:(^(x...) = OverflowContexts.saturating_pow(x...)))
+        (@__MODULE__).eval(:(abs(x) = OverflowContexts.saturating_abs(x)))
+        (@__MODULE__).eval(:(__OverflowContextDefaultSet = true))
+        nothing
+    end
+end
+
+"""
+    @checked expr
+
+Perform all integer operations in `expr` using overflow-checked arithmetic.
+"""
+macro checked(expr)
+    isa(expr, Expr) || return expr
+    expr = copy(expr)
+    return esc(replace_op!(expr, op_checked))
+end
+
 """
     @unchecked expr
 
@@ -62,14 +97,14 @@ macro unchecked(expr)
 end
 
 """
-    @checked expr
+    @saturating expr
 
-Perform all integer operations in `expr` using overflow-checked arithmetic.
+Perform all integer operations in `expr` using saturating arithmetic.
 """
-macro checked(expr)
+macro saturating(expr)
     isa(expr, Expr) || return expr
     expr = copy(expr)
-    return esc(replace_op!(expr, op_checked))
+    return esc(replace_op!(expr, op_saturating))
 end
 
 const op_checked = Dict(
@@ -92,6 +127,16 @@ const op_unchecked = Dict(
     :abs => :(unchecked_abs)
 )
 
+const op_saturating = Dict(
+    Symbol("unary-") => :(saturating_neg),
+    Symbol("ambig-") => :(saturating_negsub),
+    :+ => :(saturating_add),
+    :- => :(saturating_sub),
+    :* => :(saturating_mul),
+    :^ => :(saturating_pow),
+    :abs => :(saturating_abs)
+)
+
 const broadcast_op_map = Dict(
     :.+ => :+,
     :.- => :-,
@@ -115,6 +160,8 @@ unchecked_negsub(x) = unchecked_neg(x)
 unchecked_negsub(x, y) = unchecked_sub(x, y)
 checked_negsub(x) = checked_neg(x)
 checked_negsub(x, y) = checked_sub(x, y)
+saturating_negsub(x) = saturating_neg(x)
+saturating_negsub(x, y) = saturating_sub(x, y)
 
 # copied from CheckedArithmetic.jl and modified it
 function replace_op!(expr::Expr, op_map::Dict)
@@ -182,7 +229,7 @@ function replace_op!(expr::Expr, op_map::Dict)
     elseif isexpr(expr, :.) # broadcast function
         op = expr.args[1]
         expr.args[1] = get(op_map, op, op)
-    elseif !isexpr(expr, :macrocall) || expr.args[1] ∉ (Symbol("@checked"), Symbol("@unchecked"))
+    elseif !isexpr(expr, :macrocall) || expr.args[1] ∉ (Symbol("@checked"), Symbol("@unchecked"), Symbol("@saturating"))
         for a in expr.args
             if isa(a, Expr)
                 replace_op!(a, op_map)