Modular constant folding

The following code (taken from #16329):

func :: Int -> IO Int
func n = foldM step 0 xs
    where
    xs = map (n+) [1,2,3]
    step acc x =
        case x `mod` 3  of
            0 -> pure acc
            1 -> pure $ acc + 1
            2 -> pure $ acc + 2

func' n = foldl step 0 xs
    where
    xs = map (n+) [1,2,3]
    step acc x =
        case x `mod` 3 of
            0 -> acc
            1 -> acc + 1
            2 -> acc + 2

is simplified to the following core code:

func n =
    case n + 1 `mod` 3 of
        0 -> case n + 2 `mod` 3 of
            0 -> case n + 3 `mod` 3 of
                0 -> pure 0
                1 -> pure 1
                2 -> pure 2
            1 -> case n + 3 `mod` 3 of
                0 -> pure 1
                1 -> pure 2
                2 -> pure 3
            2 -> case n + 3 `mod` 3 of
                0 -> pure 2
                1 -> pure 3
                2 -> pure 4
        1 -> case n + 2 `mod` 3 of
            0 -> case n + 3 `mod` 3 of
                0 -> pure 1
                1 -> pure 2
                2 -> pure 3
            1 -> case n + 3 `mod` 3 of
                0 -> pure 2
                1 -> pure 3
                2 -> pure 4
            2 -> case n + 3 `mod` 3 of
                0 -> pure 3
                1 -> pure 4
                2 -> pure 5
        2 -> case n + 2 `mod` 3 of
            0 -> case n + 3 `mod` 3 of
                0 -> pure 2
                1 -> pure 3
                2 -> pure 4
            1 -> case n + 3 `mod` 3 of
                0 -> pure 3
                1 -> pure 4
                2 -> pure 5
            2 -> case n + 3 `mod` 3 of
                0 -> pure 4
                1 -> pure 5
                2 -> pure 6


func' n =
    join j2 w2 =
        join j1 w1 =
            case n + 3  `mod` 3 of
                0 -> w1
                1 -> w1 + 1
                2 -> w1 + 2
        in case n + 2 `mod` 3 of
            0 -> jump j1 w2
            1 -> jump j1 (w2 + 1)
            2 -> jump j1 (w2 + 2)
    in case n + 1 `mod` 3 of
        0 -> jump j2 0
        1 -> jump j2 1
        2 -> jump j2 2

Case-folding with modular arithmetic should remove the nesting.

changed weight to 5

added Tfeature request Trac import labels

assigned to @hsyl20

changed the description

changed the description

added Pnormal label

unassigned @hsyl20

Would the following rule be valid:

-- k and m are literals
case (n + k) `mod` m of s
  DEFAULT -> e
  l1  -> e1
  ... -> e2

===> Remove all the unreachable alternatives (lx `mod` m /= lx)

===> Then float the (+k) in:

case n `mod` m of s'
  DEFAULT -> let s = ... in e
  -- rotate all the other alternatives (we could filter unreachable ones here too):
  -- if we don't filter, an unreachable alternative may become reachable and vice-versa
  (l1 - k) `mod` m  -> let s = ... in e1
  ... -> ...

We need to compute the new scrutinee from s' in each alternative:

let s = (s' + k) `mod` m in ...

• s' is a literal in all alternatives so we can compute it statically.
• in DEFAULT alternative ideally we should avoid calling mod again (for performance). By statically computing k' = mod k m, I wonder if it is easily possible to determine using additions, subtractions, and comparisons if we should add/subtract m to s' + k' to get the proper result.

That transformation looks plausible to me. Does it enable something important?

For finite-width types, the addition n + k can wrap, causing trouble unless m (or its negation) is a power of two.