Modular arithmetic

Question

As the name says, LongModulus is a pretty long (lengthy) class implementing modular arithmetic for a long modulus. As overflow is a problem, I implemented two subclasses, one for moduli not using the highest non-sign bit (i.e., smaller than \$2^{62}\$) and one for the others.

The code uses Guava and Lombok. As usual, feel free to ignore my slightly deviating coding conventions.

My concerns are speed and correctness. I've already written a thorough test (not given here as the question is already very long), but haven't benchmarked anything yet.

The TODO markers shows the places I'd like to improve (or features currently missing like pow with negative exponent, but this is not a bug as I haven't needed it yet).

IntModulus

import static com.google.common.base.Preconditions.checkArgument;
import lombok.Getter;

/**
 * This class provides common modular arithmetic.
 * Results of all methods are ints guaranteed to be non-negative and less then modulus.
 */
public final class IntModulus {
    private IntModulus(int modulus) {
        this.modulus = modulus;
    }

    public static IntModulus newModulus(int modulus) {
        checkArgument(modulus>0);
        return new IntModulus(modulus);
    }

    public int pow(long base, long exp) {
        checkArgument(exp>=0);  //TODO allow negative exponents
        if (modulus==1) return 0;
        if (exp==0) return 1;
        return powInternal(mod(base), exp);
    }

    private int powInternal(int base, long exp) {
        assert base>=0;
        if (base<=1) return base;
        int result = 1;
        for (int x=base; exp>0; exp>>=1) {
            if ((exp&1) != 0) result = mul(result, x);
            x = square(x);
        }
        return result;
    }

    private int square(int x) {
        return mul(x, x);
    }

    public int mul(long x, long y) {
        return mul(mod(x), mod(y));
    }

    public int mul(int x, int y) {
        return mod((long) x * y);
    }

    public int add(int x, int y) {
        return mod(x + y);
    }

    public int sub(int x, int y) {
        return mod(x - y);
    }

    public int mod(long x) {
        return fixMod((int) (x % modulus));
    }

    public int mod(int x) {
        return fixMod(x % modulus);
    }

    private int fixMod(int result) {
        return (result<0 ? result+modulus : result);
    }

    @Getter private final int modulus;
}

LongModulus

import static com.google.common.base.Preconditions.checkArgument;
import javax.annotation.Nullable;
import lombok.Getter;

/**
 * This class provides common modular arithmetic.
 * Results of all methods are longs guaranteed to be non-negative and less then modulus.
 */
public abstract class LongModulus {
    private static final class HugeModulus extends LongModulus {
        HugeModulus(long modulus) {
            super(modulus);
        }

        @Override public final long pow(long base, long exp) {
            checkArgument(exp>=0);  //TODO allow negative exponents
            if (exp==0) return 1;
            base = mod(base);
            assert base>=0;
            if (base<=1) return base;
            long result = 1;
            for (long x=base; exp>0; exp>>=1) {
                if ((exp&1) != 0) result = mul(result, x);
                x = square(x);
            }
            return result;
        }
        @Override public long mul(long x, long y) {
            final long modulus = super.modulus;
            // The modulus is so damn huge that addition of reduced numbers may overflow long,
            // so we have to pretend we're using unsigned long.
            // OTOH, we need no modulus operations as a single subtraction is enough.

            final boolean negate = (x^y) < 0;
            x = Math.abs(x);
            y = Math.abs(y);

            // Handle Long.MIN_VALUE.
            if (x<0) x = super.negatedLongMinValueModulus;
            if (y<0) y = super.negatedLongMinValueModulus;

            final long x0 = low(x);
            final long x1 = high(x);
            final long y0 = low(y);
            final long y1 = high(y);

            long result = mulUnsignedInt(x1, y1);
            result = shift32(result);

            result += mulUnsignedInt(x0, y1);
            if (!isReduced(result)) result -= modulus;
            result += mulUnsignedInt(x1, y0);
            if (!isReduced(result)) result -= modulus;
            result = shift32(result);

            result += mulUnsignedInt(x0, y0);
            if (!isReduced(result)) result -= modulus;

            return negate(result, negate);
        }

        private long mulUnsignedInt(long x, long y) {
            assert isUnsignedInt(x);
            assert isUnsignedInt(y);
            final long modulus = super.modulus;
            long result = x * y;
            for (int i=0; i<3 && !isReduced(result); ++i) result -= modulus;
            assert isReduced(result) : result + " " + modulus;
            return result;
        }

        private long shift32(long x) {
            final long modulus = super.modulus;
            assert isReduced(x);
            for (int i=0; i<32; ++i) {
                x <<= 1;
                if (!isReduced(x)) x -= modulus;
                assert isReduced(x);
            }
            return x;
        }
    }

    private static final class NormalModulus extends LongModulus {
        NormalModulus(long modulus) {
            super(modulus);
            intModulus = (int) modulus == modulus ? IntModulus.newModulus((int) modulus) : null;
        }

        @Override public final long pow(long base, long exp) {
            if (intModulus!=null) return intModulus.pow(base, exp);
            checkArgument(exp>=0);  //TODO allow negative exponents
            if (exp==0) return 1;
            base = mod(base);
            assert base>=0;
            if (base<=1) return base;
            long result = 1;
            for (long x=base; exp>0; exp>>=1) {
                if ((exp&1) != 0) result = mul(result, x);
                x = square(x);
            }
            return result;
        }

        @Override public long mul(long x, long y) {
            if (intModulus!=null) return intModulus.mul(x, y);
            final long modulus = super.modulus;

            final boolean negate = (x^y) < 0;
            x = Math.abs(x);
            y = Math.abs(y);

            // Handle Long.MIN_VALUE.
            if (x<0) x = super.negatedLongMinValueModulus;
            if (y<0) y = super.negatedLongMinValueModulus;

            final long x0 = low(x);
            final long x1 = high(x);
            final long y0 = low(y);
            final long y1 = high(y);

            long result = mulUnsignedInt(x1, y1);
            result = shift32(result);

            result += mulUnsignedInt(x0, y1);
            result %= modulus;
            result += mulUnsignedInt(x1, y0);
            result %= modulus;
            result = shift32(result);

            final long d = mulUnsignedInt(x0, y0);
            result += d;
            result %= modulus;

            return negate(result, negate);
        }

        private long mulUnsignedInt(long x, long y) {
            assert isUnsignedInt(x);
            assert isUnsignedInt(y) : y + " " + super.modulus;
            final long modulus = super.modulus;
            long xy = x * y;
            if (xy>=0) return xy % modulus;

            // Handle overflow into the sign bit.
            final long lsb = xy & 1;
            xy >>>= 1;

            xy %= modulus;
            xy += xy + lsb;
            assert 0 <= xy && xy < 2*modulus;
            if (xy>modulus) xy -= modulus;
            return xy;
        }

        private long shift32(long x) {
            final long modulus = super.modulus;
            assert isReduced(x);

            for (int i=0; i<32; ++i) {
                x <<= 1;
                x %= modulus;
            }
            assert isReduced(x);
            return x;
        }

        /**
         * A delegate used in case the modulus fits in an int.
         * A specialized subclass could make method calls megamorphic, so it's avoided.
         */
        @Nullable private final IntModulus intModulus;
    }

    private LongModulus(long modulus) {
        this.modulus = modulus;
        shift32Multiplier = (1L << 32) % modulus;
        negatedLongMinValueModulus = negate(mod(Long.MIN_VALUE), true);
    }

    public static LongModulus newModulus(long modulus) {
        checkArgument(modulus>0);
        return modulus >= MIN_HUGE ? new HugeModulus(modulus) : new NormalModulus(modulus);
    }

    public abstract long pow(long base, long exp);

    final long square(long x) { //TODO optimize?
        return mul(x, x);
    }

    public abstract long mul(long x, long y);

    public final long add(long x, long y) {
        x = mod(x);
        y = mod(y);
        final long result = x + y;
        // Also handles overflow which can happen in case of huge modulus.
        return isReduced(result) ? result : result-modulus;
    }

    public final long sub(long x, long y) {
        x = mod(x);
        y = mod(y);
        return fixMod(x - y);
    }

    public final long mod(long x) {
        return fixMod(x % modulus);
    }

    private long fixMod(long result) {
        return (result<0 ? result+modulus : result);
    }

    private static long high(long x) {
        return x >> 32;
    }

    private static long low(long x) {
        return x & INT_MASK;
    }

    private static boolean isUnsignedInt(long x) {
        return 0 <= x && x <= INT_MASK;
    }

    boolean isReduced(long x) {
        return 0 <= x && x < modulus;
    }

    long negate(long x, boolean negate) {
        return !negate | x==0 ? x : modulus-x;
    }

    /** The least huge (i.e. using the most significant bit) modulus. */
    private static final long MIN_HUGE = 1L << 62;
    private static final long INT_MASK = (1L << 32) - 1;

    @Getter private final long modulus;

    private final long negatedLongMinValueModulus;

    @SuppressWarnings("unused") //TODO There should be a way to replace the slow shift by some multiplication.
    private final long shift32Multiplier;
}

UPDATE

As this question was rather hard and long, I added comments and posted a clearer followup question. More to come.

Answer 1

Well, I took a good, thorough look, and there's really just two pieces of advice I can give:

1. Comments. Use them. For example, I saw this:

   for (int x=base; exp>0; exp>>=1) {
       if ((exp&1) != 0) {
           result = this.mul(result, x);
       }
       x = this.square(x);
   }
   

   What? Maybe I'm just being dumb, but I can't figure this out for the life of me. A comment with a brief description of what's going on would help a lot -- a one-line comment at the beginning of the for stating what it does and how that ties in to the method would be enough.

2. Why are your method names so short? You use mul instead of multiply or times, and (Nevermind! See Note 1) shift32 doesn't say anything about in what direction it shifts, or its purpose. Admittedly, the purpose is probably too long for a method name, but this is a case where a Javadoc comment would be a good idea.

3. asserts -- be careful with them, especially in a library. If you use them, use the assert condition : "detail message"; form (which I do see scattered around here and there), so that if/when they throw an error, you have some details about what happened, rather than going off solely a stacktrace. Plus, it helps if you explain why you're making that assertion. I'd give any examples if I could understand your code at all (See #1 again).

Once you've done that, please do post another question. I'd love to review this code when I can actually understand it.

^{1: As maaartinus said, "long names lead to long lines pretty soon and there aren't many things as bad as broken lines."}

Answer 2

Assert statements

They're awesome. I should use them more myself, but they're not always appropriate. Let's start with when they're appropriate and then look at how you've used them.

Assert statements are fantastic for internally checking conditions during development. They get compiled away in our release build, so use them liberally anywhere we need to check a precondition... internally, in places where we have direct control on the input and if conditions don't meet our expectations, something is wrong that the dev needs to look into right now. When I say internal, I really mean internal. I'm not up to speed with Java's scoping keywords, but in C# this would be in private, internal, and (maybe) protected methods. Anytime we're in a public method, we need to be a bit more robust with our argument checking, because we have zero control over what the client hands us.

So, back to how you used them.

    @Override public final long pow(long base, long exp) {
        checkArgument(exp>=0);  //TODO allow negative exponents
        if (exp==0) return 1;
        base = mod(base);
        assert base>=0;
        if (base<=1) return base;
        long result = 1;
        for (long x=base; exp>0; exp>>=1) {

The method is public and you've done some proper argument checking. (At least I'm going to assume that returning the base is appropriate if it's < 1. Is it?) So, that's good! However, does it really make sense to Assert on a condition that has been gracefully handled? I don't think it does. Minimally, it can cause confusion for someone coming behind you wondering "why are we asserting this in s public method when we've handled that case right there.

Like I said, asserts are awesome, but try to keep them in places where there're likely to catch actual bugs.