Derive an optimized case-of expression from a function that takes a closed set as its input

I have a closed set of values:

data Value = A | B | C | D | E ...
  deriving (Eq, Ord, Show)

And a data structure that represents their order:

order :: [[Value]]
order = [
  [ B ],
  [ A, D ],
  [ C ],
  ...
  ]

I need to convert a Value's order into an Int. I could do it like this:

prec' :: [[Value]] -> Value -> Int
prec' [] _ = 0
prec' (vs : rest) v = if v `elem` vs
  then 1 + length rest
  else prec' rest v

prec :: Value -> Int
prec = prec' order

However this prec has complexity O(n).

What I would want, is a very lightweight and optimized function like this one:

prec :: Value -> Int
prec = \case
  A -> 2
  B -> 3
  C -> 1
  D -> 2
  E -> 0
  ...

But of course I don't want to write it manually, otherwise it risks being inconsistent with the information stored in order. The Haskell compiler should be able to derive that function on its own easily, since its input is a closed set.

How can I get GHC to generate a function like the latest definition of prec?

Solution 1: Use Template Haskell to generate the code you want.

Solution 2 (expanded below): (Ab)use the simplifier.

The main obstacle to simplification is that GHC will not inline recursive functions. One workaround is to do the recursion through type classes.

-- Intuitively unroll :: Nat -> (a -> a) -> (a -> a)
-- but the nat is now a type-level parameter.
class Unroll (n :: Nat) where
unroll :: (a -> a) -> (a -> a)
instance Unroll 0 where
unroll = id
instance {-# OVERLAPPABLE #-} Unroll (n-1) => Unroll n where
unroll f = f . unroll @(n-1) f

This lets you define the following fixpoint operator that unfolds the first n iterations:

unrollfix :: forall n a. Unroll n => (a -> a) -> a
unrollfix f = unroll @n f (fix f)

You then need to write all recursive functions using fix, and replace fix with unrollfix. You have to sprinkle some INLINE pragmas around too.

elem with fix:

elem :: forall a. Eq a => a -> [a] -> Bool
elem = fix go
where
go elem_ x [] = False
go elem_ x (y : ys) = x == y || elem_ x ys

elem with unrollfix:

{-# INLINE uelem #-}
uelem :: forall n a. (Unroll n, Eq a) => a -> [a] -> Bool
uelem = unrollfix @n go
where
go elem_ x [] = False
go elem_ x (y : ys) = x == y || elem_ x ys

Also length (omitted), and prec'.

prec' with fix:

prec' :: forall a. Eq a => [[a]] -> a -> Int
prec' = fix go
where
go prec_ [] v = 0
go prec_ (vs : rest) v = if elem v vs
then 1 + length rest
else prec_ rest v

prec' with unrollfix:

prec' :: forall n a. (Unroll n, Eq a) => [[a]] -> a -> Int
prec' = unrollfix @n go
where
go prec_ [] v = 0
go prec_ (vs : rest) v = if uelem @n v vs
then 1 + ulength @n rest
else prec_ rest v
{-# INLINE go #-}

Finally, set the n parameter to a high enough value to enable simplification.

prec :: Value -> Int
prec v = prec' @5 order v

Full code:

{-# LANGUAGE AllowAmbiguousTypes, DataKinds, MultiParamTypeClasses, ScopedTypeVariables, TypeApplications, UndecidableInstances #-}
{-# OPTIONS_GHC -ddump-simpl #-}
module A (Value(..), prec) where
import GHC.TypeNats
import Data.Function (fix)
import GHC.Exts
class Unroll (n :: Nat) where
unroll :: (a -> a) -> (a -> a)
instance Unroll 0 where
unroll = id
instance {-# OVERLAPPABLE #-} Unroll (n-1) => Unroll n where
unroll f = f . unroll @(n-1) f
unrollfix :: forall n a. Unroll n => (a -> a) -> a
unrollfix f = unroll @n f (fix f)
data Value = A | B | C | D | E
deriving Eq
order :: [[Value]]
order = [[A], [B, C], [D], [E]] 
{-# INLINE uelem #-}
uelem :: forall n a. (Unroll n, Eq a) => a -> [a] -> Bool
uelem = unrollfix @n go
where
go elem_ x [] = False
go elem_ x (y : ys) = x == y || elem_ x ys
{-# INLINE go #-}
{-# INLINE ulength #-}
ulength :: forall n a. Unroll n => [a] -> Int
ulength = unrollfix @n go
where
go length_ [] = 0
go length_ (_ : xs) = 1 + length_ xs
{-# INLINE go #-}
prec' :: forall n a. (Unroll n, Eq a) => [[a]] -> a -> Int
prec' = unrollfix @n go
where
{-# INLINE go #-}
go prec_ [] v = 0
go prec_ (vs : rest) v = if uelem @n v vs
then 1 + ulength @n rest
else prec_ rest v
prec :: Value -> Int
prec v = prec' @5 order v

Generated Core (using the -ddump-simpl option) (look at the unfolding, instead of the main definition):

\ (v_aQC [Occ=Once1!] :: Value) ->
case v_aQC of {
__DEFAULT -> GHC.Types.I# 3#;
A -> GHC.Types.I# 4#;
D -> GHC.Types.I# 2#;
E -> GHC.Types.I# 1#
}

I would just wrap a general-purpose memoization tool around the function, like MemoTrie or memoize or fastmemo.

{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}
import Data.Function.FastMemo
data Value = A | B | C | D | E ...
deriving (Eq, Ord, Show, Generic, Memoizable)
prec :: Value -> Int
prec = memoize $ prec' order

This may not be as optimized as a direct TH solution, but the Generic-derived Memoizable instance should result something reasonably similar. Not sure which of these packages do it best.

Define prec first, then generate order using prec.

prec :: Value -> Int
prec = \case
A -> 2
B -> 3
C -> 1
D -> 2
E -> 0
order :: [[Value]]
order = go [A, B, C, D, E]
where eqPrec = (==) `on` prec
ordPrec = compare `on` prec
go = reverse . groupBy eqPrec . sortBy ordPrec

Perhaps one simple solution would be to do by hand what you want the compiler to do automatically once it has the case statement you describe -- compute a jump table.

import Data.Array
-- deriving Enum makes a compiler-written case statement like what you want
data Foo = A | B | C | D | E deriving (Bounded, Enum)
orderArray :: Array Int Int
orderArray = listArray
(0, fromEnum (maxBound :: Foo) - 1)
(orderSlow <$> [minBound..maxBound])
prec :: Foo -> Int
prec = unsafeAt orderArray . fromEnum

This will have to run orderSlow once for each possible value, but second accesses will be fast O(1) lookups.

Rewrite your prec' function using a fold, in this case foldl', and add INLINE pragmas to prec' and order.

import Data.Foldable
import Data.Maybe

data Value = A | B | C | D | E
  deriving (Eq)

prec' :: Value -> [[Value]] -> Int
prec' v = fromMaybe 0 . foldl' f Nothing
  where
    f (Just !len) _ = Just (len + 1)
    f Nothing vs | v `elem` vs = Just 1
                 | otherwise = Nothing
{-# INLINE prec' #-}

order :: [[Value]]
order = [
  [ B ],
  [ A, D ],
  [ C ]
  ]
{-# INLINE order #-}
  
prec :: Value -> Int
prec v = prec' v order

Compiling this with GHC 9.6 I get the following simplifier output (using -ddump-simpl and some -dsuppress-* options):

prec1 :: Int
prec1 = I# 0#

prec2 :: Int
prec2 = I# 1#

-- snip --

lvl1 :: Int
lvl1 = I# 2#

lvl2 :: Int
lvl2 = I# 3#

prec :: Value -> Int
prec
  = \ (v :: Value) ->
      case v of {
        __DEFAULT -> lvl1;
        B -> lvl2;
        C -> prec2;
        E -> prec1
      }

Why does this work?

When you write a list literal, GHC will actually desugar it using the build function, which means that if you use a fold then your function can potentially participate in list fusion. For example, if I compile your order with -ddump-ds (which dumps the desugaring output as GHC Core) then I get the following result:

order
  = build
      (\ (@a_d1Pf)
         (c_d1Pg :: [Value] -> a_d1Pf -> a_d1Pf)
         (n_d1Ph :: a_d1Pf) ->
         c_d1Pg
           (build
              (\ (@a_d1P4)
                 (c_d1P5 :: Value -> a_d1P4 -> a_d1P4)
                 (n_d1P6 :: a_d1P4) ->
                 c_d1P5 B n_d1P6))
           (c_d1Pg
              (build
                 (\ (@a_d1P9)
                    (c_d1Pa :: Value -> a_d1P9 -> a_d1P9)
                    (n_d1Pb :: a_d1P9) ->
                    c_d1Pa A (c_d1Pa D n_d1Pb)))
              (c_d1Pg
                 (build
                    (\ (@a_d1Pc)
                       (c_d1Pd :: Value -> a_d1Pc -> a_d1Pc)
                       (n_d1Pe :: a_d1Pc) ->
                       c_d1Pd C n_d1Pe))
                 n_d1Ph)))

So with enough inlining and the fact that both elem and foldl' participate in list fusion as good consumers, the simplifier can optimize your function aggressively using rewrite rules.

For completeness, here is a TH solution (template-haskell-2.19.0):

{-# LANGUAGE TemplateHaskell #-}
module PrecTH where
import Language.Haskell.TH
import Data.List (nub)
prec' :: Ord a => [[a]] -> a -> Int
prec' [] _ = 0
prec' (vs : rest) v = if v `elem` vs
then 1 + length rest
else prec' rest v
mkPrecValueDataType :: String -> [[String]] -> DecsQ
mkPrecValueDataType dtName order = pure [
DataD [] dtName' [] Nothing
[ NormalC (mkName c) []
| c <- cstrs ]
[DerivClause Nothing $ ConT <$> [''Eq, ''Ord, ''Show ]]
, SigD (mkName "prec")
$ ArrowT `AppT` ConT dtName' `AppT` ConT ''Int
, FunD (mkName "prec")
[ Clause [ConP (mkName c) [] []]
(NormalB . LitE . IntegerL $ slowPrec c)
[]
| c <- cstrs ]
]
where slowPrec = fromIntegral . prec' order
cstrs = concat order -- apply `sort` here if you like the constructors in alphabetical order
dtName' = mkName dtName

To be used thus

{-# LANGUAGE TemplateHaskell #-}
module PrecValues where
import PrecTH
mkPrecValueDataType "Value" [["B"], ["A","D"], ["C"]]

and producing

ghci> :browse
prec :: Value -> Int
type Value :: *
data Value = B | A | D | C
ghci> prec<$>[A,B,C,D]
[2,3,1,2]