融合长度索引链表

Question

我有标准Vect类型:

module Vect where

data Nat = Z | S Nat
data Vect (n :: Nat) (a :: Type) where
  VNil :: Vect Z a
  VCons :: a -> Vect n a -> Vect (S n) a

我有这些功能:

foldVect :: forall (p :: Nat -> Type) n a.
            p Z ->
            (forall m. a -> p m -> p (S m)) ->
            Vect n a -> p n
foldVect n c = go
  where go :: forall l. Vect l a -> p l
        go VNil = n
        go (VCons x xs) = x `c` go xs

newtype FVect a n = FVect { unFVect :: Vect n a }
buildVect :: forall n a.
             (forall (p :: Nat -> Type).
              p Z ->
              (forall m. a -> p m -> p (S m)) ->
              p n
             ) -> Vect n a
buildVect f = unFVect $ f (FVect VNil) $ \x (FVect xs) -> FVect $ x `VCons` xs

我尝试重新创建(部分)机器base,允许列表融合:

instance Functor (Vect n) where
    fmap = mapVect
    {-# INLINE fmap #-}

mapVect :: forall n a b. (a -> b) -> (Vect n a -> Vect n b)
mapVect _f VNil = VNil
mapVect f (VCons x xs) = f x `VCons` mapVect f xs
mapFB :: forall (p :: Nat -> Type) n a b. (forall m. b -> p m -> p (S m)) -> (a -> b) -> a -> p n -> p (S n)
mapFB cons f = \x ys -> cons (f x) ys
{-# INLINE [0] mapFB #-}
{-# NOINLINE [0] mapVect #-}
{-# RULES "mapVect" [~1] forall f xs. mapVect f xs = buildVect (\nil cons -> foldVect nil (mapFB cons f) xs) #-}

{-# INLINE [0] foldVect #-}
-- base has this; I don't think it does anything without a "refolding" rule on mapVect
{-# INLINE [0] buildVect #-}
{-# RULES "foldVect/buildVect" forall (nil :: p Z)
                                      (cons :: forall m. a -> p m -> p (S m))
                                      (f :: forall (q :: Nat -> Type).
                                            q Z ->
                                            (forall m. a -> q m -> q (S m)) ->
                                            q n
                                      ).
          foldVect nil cons (buildVect f) = f nil cons
  #-}

然后

module Test where
import Vect

test :: Vect n Int -> Vect n Int
test = fmap (*5) . fmap (+2)

没有融合发生.如果我通过-ddump-simpl,我看到foldVect并且buildVect已经内联,但......

1. 该INLINEs的阶段性,使他们不干扰融合反正.(毕竟,这是怎么base会为[])

2. INLINE [0]用NOINLINE油漆替换s是一个相当惊人的图像:

   test
     = \ (@ (n_a141 :: Nat)) (x_X1lK :: Vect n_a141 Int) ->
         buildVect
           @ n_a141
           @ Int
           (\ (@ (p_X1jl :: Nat -> *))
              (nil_X11K [OS=OneShot] :: p_X1jl 'Z)
              (cons_X11M [OS=OneShot]
                 :: forall (m :: Nat). Int -> p_X1jl m -> p_X1jl ('S m)) ->
              foldVect
                @ p_X1jl
                @ n_a141
                @ Int
                nil_X11K
                (\ (@ (m_a1i5 :: Nat))
                   (x1_aYI :: Int)
                   (ys_aYJ [OS=OneShot] :: p_X1jl m_a1i5) ->
                   cons_X11M
                     @ m_a1i5
                     (case x1_aYI of { GHC.Types.I# x2_a1l5 ->
                      GHC.Types.I# (GHC.Prim.*# x2_a1l5 5#)
                      })
                     ys_aYJ)
                (buildVect
                   @ n_a141
                   @ Int
                   (\ (@ (p1_a1i0 :: Nat -> *))
                      (nil1_a10o [OS=OneShot] :: p1_a1i0 'Z)
                      (cons1_a10p [OS=OneShot]
                         :: forall (m :: Nat). Int -> p1_a1i0 m -> p1_a1i0 ('S m)) ->
                      foldVect
                        @ p1_a1i0
                        @ n_a141
                        @ Int
                        nil1_a10o
                        (\ (@ (m_a1i5 :: Nat))
                           (x1_aYI :: Int)
                           (ys_aYJ [OS=OneShot] :: p1_a1i0 m_a1i5) ->
                           cons1_a10p
                             @ m_a1i5
                             (case x1_aYI of { GHC.Types.I# x2_a1lh ->
                              GHC.Types.I# (GHC.Prim.+# x2_a1lh 2#)
                              })
                             ys_aYJ)
                        x_X1lK)))
   

   一切都在那里,但简化器就是没有它.

如果我检查规则本身,我会看到这一点

"foldVect/buildVect"
    forall (@ (p_aYG :: Nat -> *))
           (@ (n_aYJ :: Nat))
           (@ a_aYH)
           (nil_aYD :: p_aYG 'Z)
           (cons_aYE :: forall (m :: Nat). a_aYH -> p_aYG m -> p_aYG ('S m))
           (f_aYF
              :: forall (q :: Nat -> *).
                 q 'Z -> (forall (m :: Nat). a_aYH -> q m -> q ('S m)) -> q n_aYJ).
      foldVect @ p_aYG
               @ n_aYJ
               @ a_aYH
               nil_aYD
               cons_aYE
               (buildVect
                  @ n_aYJ
                  @ a_aYH
                  (\ (@ (p1_a156 :: Nat -> *))
                     (ds_d1io :: p1_a156 'Z)
                     (ds1_d1ip
                        :: forall (m :: Nat). a_aYH -> p1_a156 m -> p1_a156 ('S m)) ->
                     f_aYF @ p1_a156 ds_d1io ds1_d1ip))
      = f_aYF @ p_aYG nil_aYD cons_aYE

似乎问题是buildVect需要成为一个非常特定形式的lambda抽象的论证,而我在构建一个最终发生的重写系统时遇到了麻烦.

我如何融合工作？

(我不知道这是否有用甚至是正确的;我只是这样做,看我是否可以.)

Answer 1

像往常一样，newtype每当编译器被胡言乱语时，就可以挽救局面：

module Vect where
-- everything else the same...
newtype VectBuilder n a = VectBuilder { runVectBuilder :: forall (p :: Nat -> Type).
                                                          p Z ->
                                                          (forall m. a -> p m -> p (S m)) ->
                                                          p n
                                      }

buildVect' :: forall n a. VectBuilder n a -> Vect n a
buildVect' f = unFVect $
                runVectBuilder f (FVect VNil) $ \x (FVect xs) -> FVect $ x `VCons` xs
{-# INLINE [0] buildVect' #-}
buildVect :: forall n a.
             (forall (p :: Nat -> Type).
              p Z ->
              (forall m. a -> p m -> p (S m)) ->
              p n
             ) -> Vect n a
buildVect f = buildVect' (VectBuilder f)
{-# INLINE buildVect #-}

{-# RULES "foldVect/buildVect'" forall (nil :: p Z)
                                       (cons :: forall m. a -> p m -> p (S m))
                                       (f :: VectBuilder n a).
                                foldVect nil cons (buildVect' f) = runVectBuilder f nil cons
  #-}
-- compiler no longer has a chance to muck up the LHS by eta expanding f because
-- f "isn't" a function anymore

-- rule for mapVect goes unchanged, so I guess that's evidence that this is totally transparent

module Test where
import Vect
test :: Vect n Int -> Vect n Int
test = fmap (*5) . fmap (+2)

module Vect where
-- everything else the same...
newtype VectBuilder n a = VectBuilder { runVectBuilder :: forall (p :: Nat -> Type).
                                                          p Z ->
                                                          (forall m. a -> p m -> p (S m)) ->
                                                          p n
                                      }

buildVect' :: forall n a. VectBuilder n a -> Vect n a
buildVect' f = unFVect $
                runVectBuilder f (FVect VNil) $ \x (FVect xs) -> FVect $ x `VCons` xs
{-# INLINE [0] buildVect' #-}
buildVect :: forall n a.
             (forall (p :: Nat -> Type).
              p Z ->
              (forall m. a -> p m -> p (S m)) ->
              p n
             ) -> Vect n a
buildVect f = buildVect' (VectBuilder f)
{-# INLINE buildVect #-}

{-# RULES "foldVect/buildVect'" forall (nil :: p Z)
                                       (cons :: forall m. a -> p m -> p (S m))
                                       (f :: VectBuilder n a).
                                foldVect nil cons (buildVect' f) = runVectBuilder f nil cons
  #-}
-- compiler no longer has a chance to muck up the LHS by eta expanding f because
-- f "isn't" a function anymore

-- rule for mapVect goes unchanged, so I guess that's evidence that this is totally transparent

这个故事的寓意是：排名足够高的类型会使RULES系统崩溃，所以给 GHC 一些关于newtypes 的帮助，即使它们在其他方面不是必要的。