SPECIALIZE is not always transitive within the same module
I recently ran into this issue while working on a re-implementation of the free var collector. According to the GHC manual:
A
SPECIALIZEhas the effect of generating (a) a specialised version of the function and (b) a rewrite rule (see :ref:rewrite-rules) that rewrites a call to the un-specialised function into a call to the specialised one. Moreover, given aSPECIALIZEpragma for a functionf, GHC will automatically create specialisations for any type-class-overloaded functions called byf, if they are in the same module as theSPECIALIZEpragma, or if they areINLINABLE; and so on, transitively.
Here's a small example where the "transitively" part isn't respected:
-- Lib.hs
module Lib
( isEven, isOdd )
where
{-# SPECIALIZE isEven :: Int -> Bool #-}
isEven :: Integral a => a -> Bool
isEven = fst evenOdd
{-# SPECIALIZE isOdd :: Int -> Bool #-}
isOdd :: Integral a => a -> Bool
isOdd = snd evenOdd
evenOdd :: Integral a => (a -> Bool, a -> Bool)
evenOdd = (goEven, goOdd)
where
goEven n
| n < 0 = goEven (- n)
| n > 0 = goOdd (n - 1)
| otherwise = True
goOdd n
| n < 0 = goOdd n
| otherwise = goEven n
-- Main.hs
{-# LANGUAGE TypeApplications #-}
module Main where
import Lib ( isEven )
main :: IO ()
main = print $ isEven @Int 10
Even though isEven is specialized, the call to isEven in main still passes around a dictionary (see attached files for the full core outputMain.dump-simpl
case Lib.$wevenOdd @Int GHC.Real.$fIntegralInt of
...
Checking the core output of Lib.hs, I can see a specialized version of isEven being produced:
Lib.isEven_goEven [InlPrag=NOUSERINLINE[2]] :: Int -> Bool
with the correct rewrite rule at the bottom:
"SPEC isEven"
forall ($dIntegral_a1Md :: Integral Int).
isEven @Int $dIntegral_a1Md
= Lib.isEven_goEven
For some reason that I don't understand, if I specialize evenOdd and isEven, then the call to isEven in Main will get specialized properly.
Note that the example first defines the evenOdd tuple and then projects the two functions out before exporting. If we remove the indirectness of evenOdd and define isEven and isOdd directly at the top-level, then the call in main module is specialized properly:
module Lib
( isEven )
where
{-# SPECIALIZE isEven :: Int -> Bool #-}
isEven :: Integral a => a -> Bool
isEven n
| n < 0 = isEven (- n)
| n > 0 = isOdd (n - 1)
| otherwise = True
isOdd :: Integral a => a -> Bool
isOdd n
| n < 0 = isEven n
| otherwise = isOdd n
Here's a link to a repo with the examples set up in a cabal project: https://gitlab.haskell.org/yiyunliu/specialization-transitivity-test
The oldest commit is the problematic case where it doesn't specialize. The rest two commits both specialize properly.