polymorphic-unboxed.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

module Test (
  geometricMean
) where

import Data.Foldable (Foldable)
import qualified Data.Foldable as F

class Tuply a where
  data Elem a :: * 
  mkTuply  :: a -> Int -> Elem a
  getSum   :: Elem a -> a
  getCount :: Elem a -> Int

instance Tuply Double where
  data Elem Double  = TD {-# UNPACK #-}!Double {-# UNPACK #-}!Int
  mkTuply a b       = TD a b
  getSum   (TD s _) = s
  getCount (TD _ c) = c
  {-# INLINE mkTuply #-}
  {-# INLINE getSum #-}
  {-# INLINE getCount #-}

  
geometricMean :: (Floating a, Tuply a, Foldable f) => f a -> a
geometricMean = fini . F.foldl' go (mkTuply 1 0)
  where
    fini t  = getSum t ** (1 / fromIntegral (getCount t))
    go t x = mkTuply (getSum t * x) (getCount t + 1) 
{-# SPECIALIZE geometricMean :: Foldable f => f Double -> Double #-}