从采用闭集作为输入的函数导出优化的 case-of 表达式

Question

我有一组封闭的值：

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)。

我想要的是一个非常轻量级且优化的函数，如下所示：

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

但我当然不想手动编写它，否则有与存储的信息不一致的风险order。Haskell 编译器应该能够轻松地自行导出该函数，因为它的输入是一个闭集。

如何让 GHC 生成类似 的最新定义的函数prec？

Answer 1

解决方案1：使用Template Haskell 生成您想要的代码。

解决方案 2（在下面进行扩展）：(Ab)使用简化器。

简化的主要障碍是 GHC 不会内联递归函数。一种解决方法是通过类型类进行递归。

-- 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

这允许您定义以下展开前 n 次迭代的定点运算符：

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

然后，您需要使用 编写所有递归函数fix，并替换fix为unrollfix。您还必须洒一些INLINE实用程序。

elem和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和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

又length（略），和prec'。

prec'和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'和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 #-}

最后，将n参数设置为足够高的值以实现简化。

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

完整代码：

{-# 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

生成的核心（使用-ddump-simpl选项）（查看展开，而不是主要定义）：

\ (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#
                 }