PartialTypeSigs.hs

{-# LANGUAGE TemplateHaskell #-}
module PartialTypeSigs where


import qualified Data.Map as M
import Data.IORef
import System.IO.Unsafe
import Data.Maybe
import Language.Haskell.TH

{-# NOINLINE m #-}
m :: IORef (M.Map String Int)
m = unsafePerformIO (newIORef M.empty)

sigs :: [ ExpQ ] -> DecsQ
sigs es = do
    runIO $ writeIORef m M.empty
    fmap concat $ mapM unifiesWith es
    
unifiesWithPrefix = "partialTypeSig_"

unifiesWith :: ExpQ -> DecsQ
unifiesWith e = do
    SigE (LitE (StringL e)) t <- e
    k <- runIO $ atomicModifyIORef m $ \k ->
          (M.insertWith (+) e 1 k, fromMaybe 1 $ M.lookup e k)
    unifiesWith1 (unifiesWithPrefix++e++show k) (dyn e) (return t)

unifiesWith1 :: String -> ExpQ -> TypeQ -> DecsQ
unifiesWith1 s e t = do
    x <- newName "x"
    fmap (:[]) $ funD (mkName s)
      [clause [varP x]
      (normalB [| ($(varE x) `asTypeOf` $e) `asTypeOf` (undefined :: $t) |])
      []]


constrainUnifiesWith :: String -> ExpQ -> ExpQ
constrainUnifiesWith k args = do
    args <- args
    m <- runIO (readIORef m)
    maybe [| error "constrainUnifiesWith" |] (toExp (unappsE args)) $ M.lookup k m
  where
    toExp :: [Exp] -> Int -> ExpQ
    toExp args n = 
      foldr (\x y -> [| $x `asTypeOf` $y |])
        [| error "constrainUnifiesWith" |]
        [ foldr (flip appE)
              [| $(dyn (unifiesWithPrefix++k++show i)) undefined |]
              (map return args)
            | i <- [1 .. n] ]


-- somewhat in reverse: [| x (f y) z |] -> [ z, f y , x]
unappsE :: Exp -> [Exp]
unappsE (AppE x y) = y : unappsE x
unappsE x = [x]

PartialTypeSigsEx.hs

{-# LANGUAGE TemplateHaskell #-}
module PartialTypeSigsEx where
import Control.Monad
import PartialTypeSigs
import Language.Haskell.TH

sigs
    [ [| "f" :: a -> b -> (a, Int) |],
      [| "f" :: b -> a -> (Char, a) |] ]

f x y | False = $(constrainUnifiesWith "f" [| x y |])
f x y = (x,y)

{- this should infer

f :: Char -> Int -> (Char,Int)

but instead "sigs" fails because it generates two
decls with the same name. IE:


PartialTypeSigsEx.hs:1:1: Splicing declarations
    sigs
      [[| "f" :: a -> b -> (a, Int) |], [| "f" :: b -> a -> (Char, a) |]]
  ======>
    PartialTypeSigsEx.hs:(7,1)-(9,40)
    partialTypeSig_f1 x
      = ((x `asTypeOf` f)
         `asTypeOf` (undefined :: forall a b. a -> b -> (a, Int)))
    partialTypeSig_f1 x
      = ((x `asTypeOf` f)
         `asTypeOf` (undefined :: forall b a. b -> a -> (Char, a)))


Shouldn't the IORef "m" get updated so that the second defn is called f2?
-}