efficiently checking that all the elements of a (big) list are the same

前端 未结 9 1033
清酒与你
清酒与你 2020-12-14 16:01

Problem

Let us suppose that we have a list xs (possibly a very big one), and we want to check that all its elements are the same.

I came up wi

9条回答
  •  时光说笑
    2020-12-14 16:46

    Here's another fun way:

    {-# INLINABLE allSame #-}
    allSame :: Eq a => [a] -> Bool
    allSame xs = foldr go (`seq` True) xs Nothing where
      go x r Nothing = r (Just x)
      go x r (Just prev) = x == prev && r (Just x)
    

    By keeping track of the previous element, rather than the first one, this implementation can easily be changed to implement increasing or decreasing. To check all of them against the first instead, you could rename prev to first, and replace Just x with Just first.


    How will this be optimized? I haven't checked in detail, but I'm going to tell a good story based on some things I know about GHC's optimizations.

    Suppose first that list fusion does not occur. Then foldr will be inlined, giving something like

    allSame xs = allSame' xs Nothing where
      allSame' [] = (`seq` True)
      allSame' (x : xs) = go x (allSame' xs)
    

    Eta expansion then yields

    allSame' [] acc = acc `seq` True
    allSame' (x : xs) acc = go x (allSame' xs) acc
    

    Inlining go,

    allSame' [] acc = acc `seq` True
    allSame' (x : xs) Nothing = allSame' xs (Just x)
    allSame' (x : xs) (Just prev) =
      x == prev && allSame' xs (Just x)
    

    Now GHC can recognize that the Maybe value is always Just on the recursive call, and use a worker-wrapper transformation to take advantage of this:

    allSame' [] acc = acc `seq` True
    allSame' (x : xs) Nothing = allSame'' xs x
    allSame' (x : xs) (Just prev) = x == prev && allSame'' xs x
    
    allSame'' [] prev = True
    allSame'' (x : xs) prev = x == prev && allSame'' xs x
    

    Remember now that

    allSame xs = allSame' xs Nothing
    

    and allSame' is no longer recursive, so it can be beta-reduced:

    allSame [] = True
    allSame (x : xs) = allSame'' xs x
    
    allSame'' [] _ = True
    allSame'' (x : xs) prev = x == prev && allSame'' xs x
    

    So the higher-order code has turned into efficient recursive code with no extra allocation.

    Compiling the module defining allSame using -O2 -ddump-simpl -dsuppress-all -dno-suppress-type-signatures yields the following (I've cleaned it up a bit):

    allSame :: forall a. Eq a => [a] -> Bool
    allSame =
      \ (@ a) ($dEq_a :: Eq a) (xs0 :: [a]) ->
        let {
          equal :: a -> a -> Bool
          equal = == $dEq_a } in
        letrec {
          go :: [a] -> a -> Bool
          go =
            \ (xs :: [a]) (prev :: a) ->
              case xs of _ {
                [] -> True;
                : y ys ->
                  case equal y prev of _ {
                    False -> False;
                    True -> go ys y
                  }
              }; } in
        case xs0 of _ {
          [] -> True;
          : x xs -> go xs x
        }
    

    As you can see, this is essentially the same as the result I described. The equal = == $dEq_a bit is where the equality method is extracted from the Eq dictionary and saved in a variable so it only needs to be extracted once.


    What if list fusion does occur? Here's a reminder of the definition:

    allSame xs = foldr go (`seq` True) xs Nothing where
      go x r Nothing = r (Just x)
      go x r (Just prev) = x == prev && r (Just x)
    

    If we call allSame (build g), the foldr will fuse with the build according to the rule foldr c n (build g) = g c n, yielding

    allSame (build g) = g go (`seq` True) Nothing
    

    That doesn't get us anywhere interesting unless g is known. So let's choose something simple:

    replicate k0 a = build $ \c n ->
      let
        rep 0 = n
        rep k = a `c` rep (k - 1)
      in rep k0
    

    So if h = allSame (replicate k0 a), h becomes

    let
      rep 0 = (`seq` True)
      rep k = go a (rep (k - 1))
    in rep k0 Nothing
    

    Eta expanding,

    let
      rep 0 acc = acc `seq` True
      rep k acc = go a (rep (k - 1)) acc
    in rep k0 Nothing
    

    Inlining go,

    let
      rep 0 acc = acc `seq` True
      rep k Nothing = rep (k - 1) (Just a)
      rep k (Just prev) = a == prev && rep (k - 1) (Just a)
    in rep k0 Nothing
    

    Again, GHC can see the recursive call is always Just, so

    let
      rep 0 acc = acc `seq` True
      rep k Nothing = rep' (k - 1) a
      rep k (Just prev) = a == prev && rep' (k - 1) a
      rep' 0 _ = True
      rep' k prev = a == prev && rep' (k - 1) a
    in rep k0 Nothing
    

    Since rep is no longer recursive, GHC can reduce it:

    let
      rep' 0 _ = True
      rep' k prev = a == prev && rep' (k - 1) a
    in
      case k0 of
        0 -> True
        _ -> rep' (k - 1) a
    

    As you can see, this can run with no allocation whatsoever! Obviously, it's a silly example, but something similar will happen in many more interesting cases. For example, if you write an AllSameTest module importing the allSame function and defining

    foo :: Int -> Bool
    foo n = allSame [0..n]
    

    and compile it as described above, you'll get the following (not cleaned up).

    $wfoo :: Int# -> Bool
    $wfoo =
      \ (ww_s1bY :: Int#) ->
        case tagToEnum# (># 0 ww_s1bY) of _ {
          False ->
            letrec {
              $sgo_s1db :: Int# -> Int# -> Bool
              $sgo_s1db =
                \ (sc_s1d9 :: Int#) (sc1_s1da :: Int#) ->
                  case tagToEnum# (==# sc_s1d9 sc1_s1da) of _ {
                    False -> False;
                    True ->
                      case tagToEnum# (==# sc_s1d9 ww_s1bY) of _ {
                        False -> $sgo_s1db (+# sc_s1d9 1) sc_s1d9;
                        True -> True
                      }
                  }; } in
            case ww_s1bY of _ {
              __DEFAULT -> $sgo_s1db 1 0;
              0 -> True
            };
          True -> True
        }
    
    foo :: Int -> Bool
    foo =
      \ (w_s1bV :: Int) ->
        case w_s1bV of _ { I# ww1_s1bY -> $wfoo ww1_s1bY }
    

    That may look disgusting, but you'll note that there are no : constructors anywhere, and that the Ints are all unboxed, so the function can run with zero allocation.

提交回复
热议问题