如何让GHC自动生成针对封闭Value类型的优化prec函数？

我有一个封闭的值集合：

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

还有一个表示它们优先级顺序的数据结构：

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

我需要将Value转换为对应的优先级Int类型。目前的实现是：

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

但这个prec函数的时间复杂度是O(n)。

我想要一个轻量且优化后的函数，类似下面这种直接模式匹配的形式（时间复杂度O(1)）：

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

但我不想手动编写这个函数，否则很容易和order中定义的优先级顺序不一致。由于Value是封闭集合，GHC理论上可以自动推导生成这个函数。请问如何让GHC生成类似上述定义的prec函数？

最直接的方式是使用Template Haskell (TH) 在编译期自动生成模式匹配形式的prec函数——既保证和order的定义严格一致，又能得到GHC优化后的O(1)分支跳转效果。

实现步骤 #

1. 定义优先级映射 #

先将order转换成键值对列表，记录每个Value对应的优先级：

priorityMap :: [(Value, Int)]
priorityMap = concatMap assignPriority $ zip [length order, length order -1 .. 1] order
  where
    assignPriority (p, vs) = map (\v -> (v, p)) vs

比如示例中的order会生成[(B,3), (A,2), (D,2), (C,1)]，未出现在order中的Value后续会默认返回0，和原prec'的行为一致。

2. 用Template Haskell生成模式匹配函数 #

导入Language.Haskell.TH模块，编写TH函数在编译期生成prec的定义：

完整代码示例：

{-# LANGUAGE TemplateHaskell #-}

module Priority where

import Language.Haskell.TH

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

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

-- 生成优先级键值对映射
priorityMap :: [(Value, Int)]
priorityMap = concatMap assignPriority $ zip [length order, length order -1 .. 1] order
  where
    assignPriority (p, vs) = map (\v -> (v, p)) vs

-- Template Haskell 生成prec函数的定义
genPrec :: Q [Dec]
genPrec = do
  -- 获取Value类型的所有构造函数
  constrInfo <- reify ''Value
  let constrs = case constrInfo of
        TyConI (DataD _ _ _ _ cs _) -> map (\(NormalC name _) -> name) cs
        _ -> error "Value必须是简单的代数数据类型"
  
  -- 为每个构造函数生成匹配分支
  branches <- mapM (\conName -> do
        let val = read (nameBase conName) :: Value
        let priority = maybe 0 id (lookup val priorityMap)
        return $ match (conP conName []) (normalB (litE (integerL (toInteger priority)))) []
      ) constrs
  
  -- 生成prec函数的完整定义
  return [FunD 'prec [Clause [] (NormalB (CaseE (VarE 'v) branches)) []]]

-- 触发编译期代码生成
$(genPrec)

3. 效果说明 #

• 编译时，Template Haskell会自动生成和你手动编写完全一致的模式匹配代码，GHC会将其优化为O(1)的跳转表，效率最高。
• 只要修改order的定义，prec函数的逻辑会自动同步，完全避免手动维护带来的不一致问题。

替代方案：基于Enum的数组索引 #

如果Value实现了Enum和Bounded，也可以用数组存储优先级，通过fromEnum索引实现O(1)访问：

import Data.Array

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

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

priorityArr :: Array Int Int
priorityArr = array (fromEnum minVal, fromEnum maxVal)
  [ (fromEnum v, prec' order v) | v <- [minVal .. maxVal] ]
  where
    minVal = minBound :: Value
    maxVal = maxBound :: Value

prec :: Value -> Int
prec v = priorityArr ! fromEnum v

这种方式代码更简洁，但生成的是数组访问逻辑，而非模式匹配。如果需要严格和目标代码形式一致，优先选择Template Haskell方案。

内容的提问来源于stack exchange，提问作者Blue Nebula