<commit_before>module FitSpec
  ( Memo (..)
  , initMemo
  , lookup
  , insert
  , toList

  , Listate (..)
  , Listate2 (..)
  , Listate3 (..)
  , runListate
  , runListate2
  , runListate3

  , lsMutateApply

  , lsMutateApply12
  , lsMutateApply22

  , lsMutateApply13
  , lsMutateApply23
  , lsMutateApply33

  , Sized (..)

  , Args (..)
  , args

  , lholds

  , report
  , reportWith
  , report2
  , report2With
  , report3
  , report3With

  , fitspec
  )
where

import Test.Check
import Control.Monad
import Data.List (transpose,sortBy)
import Data.Ord (comparing)
import Data.Maybe (isJust, isNothing, catMaybes, listToMaybe)
import Table
import qualified Data.Map as Map (insert,lookup,empty,Map,toList)
import Prelude hiding (lookup)
import Control.Applicative (liftA, liftA2)
import Data.Functor.Compose
import Utils


-- In the end, this can be merged with Listable.  It can be even generated by
-- the same combinators: in a similar way to what Uniplate does when generating
-- its different traversing functions
class Sized a where
  size :: a -> Int

instance Sized Int where
  size = id

instance Sized Bool where
  size _ = 0

instance (Sized a) => Sized [a] where
  size = sum . map ((+1) . size)

instance (Sized a, Sized b) => Sized (a,b) where
  size (x,y) = size x + size y


data Memo a b = Memo (a->b) (Map.Map a (Maybe b))

initMemo :: (a->b) -> (Memo a b)
initMemo f = Memo f Map.empty

lookup :: Ord a => a -> Memo a b -> Maybe (Maybe b)
lookup x (Memo f bs) = Map.lookup x bs

insert :: Ord a => a -> Maybe b -> Memo a b -> Memo a b
insert x y (Memo f bs) = Memo f (Map.insert x y bs)

toList :: Memo a b -> [(a,b)]
toList (Memo _ bs) = map (\(x,Just y) -> (x,y))
                   . filter (\(_,m) -> isJust m)
                   $ Map.toList bs

instance (Eq b, Show a, Show b) => Show (Memo a b) where
  show m = ("Memo" ++) . show $ toList m

data Listate a b c = Listate { listates :: Memo a b
                                        -> [[ (c,Memo a b) ]] -- a listing
                             }

data Listate2 a b c d e = Listate2 { listates2 :: Memo a b
                                               -> Memo c d
                                               -> [[ (e,Memo a b,Memo c d) ]] }

data Listate3 a b c d e f g = Listate3 { listates3 :: Memo a b
                                                   -> Memo c d
                                                   -> Memo e f
                                                   -> [[ (g,Memo a b,Memo c d,Memo e f) ]] }

runListate :: (a->b)
           -> Listate a b c
           -> [[ (c,Memo a b) ]]
runListate f = (`listates` initMemo f)


runListate2 :: (a->b)
            -> (c->d)
            -> Listate2 a b c d e
            -> [[ (e,Memo a b,Memo c d) ]]
runListate2 f g = \ls -> listates2 ls (initMemo f) (initMemo g)


runListate3 :: (a->b)
            -> (c->d)
            -> (e->f)
            -> Listate3 a b c d e f g
            -> [[ (g,Memo a b,Memo c d,Memo e f) ]]
runListate3 f g h = \ls -> listates3 ls (initMemo f) (initMemo g) (initMemo h)


lsMutateApply :: (Ord a, Eq b, Sized a, Listable b)
              => a
              -> Listate a b b
lsMutateApply x = Listate $ lsMutateApply' (,) x


lsMutateApply' :: (Eq b, Ord a, Listable b, Sized a)
               => (b -> Memo a b -> t)
               -> a
               -> Memo a b
               -> [[t]]
lsMutateApply' tup x (m@(Memo f bs)) =
  case lookup x m of
    Nothing  -> [f x `tup` insert x Nothing m]
              : case errorToNothing (f x) of
                  Just (fx) -> replicate (size x) []
                            ++ lsmap (\y -> y `tup` insert x (Just y) m) (lsfilter (/= fx) listing)
                  Nothing   -> []  -- we don't mutate undefined / error values
    (Just Nothing)  -> [[ f x `tup` m ]]
    (Just (Just y)) -> [[ y `tup` m ]]


lsMutateApply12 :: (Ord a, Eq b, Sized a, Listable b)
                => a
                -> Listate2 a b c d b
lsMutateApply12 x = Listate2 $ \m m' -> lsMutateApply' (\x m'' -> (x,m'',m')) x m

lsMutateApply22 :: (Ord c, Eq d, Sized c, Listable d)
                => c
                -> Listate2 a b c d d
lsMutateApply22 x = Listate2 $ \m' m -> lsMutateApply' (\x m'' -> (x,m',m'')) x m


lsMutateApply13 :: (Ord a, Eq b, Sized a, Listable b)
                => a
                -> Listate3 a b c d e f b
lsMutateApply13 x = Listate3 $ \m n o -> lsMutateApply' (\x m' -> (x,m',n,o)) x m

lsMutateApply23 :: (Ord c, Eq d, Sized c, Listable d)
                => c
                -> Listate3 a b c d e f d
lsMutateApply23 x = Listate3 $ \m n o -> lsMutateApply' (\x n' -> (x,m,n',o)) x n

lsMutateApply33 :: (Ord e, Eq f, Sized e, Listable f)
                => e
                -> Listate3 a b c d e f f
lsMutateApply33 x = Listate3 $ \m n o -> lsMutateApply' (\x o' -> (x,m,n,o')) x o


instance Functor (Listate a b) where
  fmap = liftM

instance Applicative (Listate a b) where
  pure = return
  (<*>) = ap

instance Monad (Listate a b) where
  return x = Listate $ \m -> [[ (x,m) ]]
  ff >>= gg = Listate $ \m -> lsConcatMap (\(x,m') -> listates (gg x) m') $ listates ff m

instance Functor (Listate2 a b c d) where
  fmap = liftM

instance Applicative (Listate2 a b c d) where
  pure = return
  (<*>) = ap

instance Monad (Listate2 a b c d) where
  return x = Listate2 $ \m m' -> [[ (x,m,m') ]]
  ff >>= gg = Listate2 $ \m m' -> lsConcatMap (\(x,m'',m''') -> listates2 (gg x) m'' m''') $ listates2 ff m m'

instance Functor (Listate3 a b c d e f) where
  fmap = liftM

instance Applicative (Listate3 a b c d e f) where
  pure = return
  (<*>) = ap

instance Monad (Listate3 a b c d e f) where
  return x = Listate3 $ \m m' m'' -> [[ (x,m,m',m'') ]]
  ff >>= gg = Listate3
            $ \m m' m'' -> lsConcatMap (\(x,m''',m'''',m''''') -> listates3 (gg x) m''' m'''' m''''')
                                     $ listates3 ff m m' m''

ltest :: Monad m
      => (c -> m Bool)
      -> [c]
      -> m Bool
ltest prop xs = foldr (liftM2 (&&)) (return True) (map prop xs)

lholds :: (Monad m, Listable c)
       => Int
       -> (c -> m Bool)
       -> m Bool
lholds n prop = ltest prop (take n list)


-- | Extra arguments for reporting functions
data Args = Args
          { functionNames :: [String]
          , variableNames :: [String]
          , limitResults :: Maybe Int -- maximum number of results
          }

-- | Default arguments
args = Args { functionNames = repeat "function"
            , variableNames = repeat "x"
            , limitResults = Nothing    -- show everything
            }

report :: (Show a, Show b, Eq b) => Int -> [[([Bool], Memo a b)]] -> IO ()
report nf rss = reportWith args nf rss

reportWith :: (Show a, Show b, Eq b) => Args -> Int -> [[([Bool], Memo a b)]] -> IO ()
reportWith args nf rss = putStrLn
                       . table "   "
                       . map showResult
                       . maybe id take (limitResults args)
                       $ fitspec fst snd nf rss
  where showResult (x,y,mm) = [ show x, show y, showM mm ]
        showM Nothing  = ""
        showM (Just m) = showAsCaseExp (functionNames args !! 0) (variableNames args !! 0) m

report2 :: (Show a, Show b, Eq b, Show c, Show d, Eq d)
        => Int -> [[([Bool], Memo a b, Memo c d)]] -> IO ()
report2 nf rss = report2With args nf rss

report2With :: (Show a, Show b, Eq b, Show c, Show d, Eq d)
            => Args -> Int -> [[([Bool], Memo a b, Memo c d)]] -> IO ()
report2With args nf rss = putStrLn
                        . table "   "
                        . map showResult
                        . maybe id take (limitResults args)
                        $ fitspec (\(x,_,_)->x) (\(_,m,n)->(m,n)) nf rss
  where showResult (x,y,mm) = [ show x, show y, showM mm ]
        showM Nothing  = ""
        showM (Just (m,n)) = showAsCaseExp (functionNames args !! 0) (variableNames args !! 0) m
                          ++ showAsCaseExp (functionNames args !! 1) (variableNames args !! 1) n

report3 :: (Show a, Show b, Eq b, Show c, Show d, Eq d, Show e, Show f, Eq f)
        => Int -> [[([Bool], Memo a b, Memo c d, Memo e f)]] -> IO ()
report3 nf rss = report3With args nf rss

report3With :: (Show a, Show b, Eq b, Show c, Show d, Eq d, Show e, Show f, Eq f)
            => Args -> Int -> [[([Bool], Memo a b, Memo c d, Memo e f)]] -> IO ()
report3With args nf rss = putStrLn
                        . table "   "
                        . map showResult
                        . maybe id take (limitResults args)
                        $ fitspec (\(x,_,_,_)->x) (\(_,m,n,o)->(m,n,o)) nf rss
  where showResult (x,y,mm) = [ show x, show y, showM mm ]
        showM Nothing  = ""
        showM (Just (m,n,o)) = showAsCaseExp (functionNames args !! 0) (variableNames args !! 0) m
                            ++ showAsCaseExp (functionNames args !! 1) (variableNames args !! 1) n
                            ++ showAsCaseExp (functionNames args !! 2) (variableNames args !! 2) o

fitspec :: (a -> [Bool]) -> (a -> b)
        -> Int -> [[a]] -> [(Int, [Int], Maybe b)]
fitspec getrs getsm n rss@((p:_):_) =
      filterU relevant
    . sortBy (comparing $ \(s,ids,_) -> (s, length ids))
    . tag
    . fitspec' getrs getsm n
    $ rss
  where tag = zipWith (\ids (ns,sms) -> (ns,ids,sms)) (subsets [1..length (getrs p)])
        (n,is,_) `relevant` (n',is',_) = not (n == n' && is `contained` is')
fitspec _ _ _ _ = error "fitspec: something went wrong, no results to analyse?"

fitspec' :: (a -> [Bool]) -> (a -> b)
         -> Int -> [[a]] -> [(Int, Maybe b)]
fitspec' getrs getsm n =
      map (\fs -> (count isJust fs, listToMaybe $ catMaybes fs))
    . transpose
    . map (\p -> map (\b -> if b then Just (getsm p) else Nothing) (compositions (getrs p)))
    . take n . tail . concat
  where count f = length . filter f


-- | 'compositions' @bs@ returns all compositions formed by taking values of @bs@
compositions :: [Bool] -> [Bool]
compositions = map and . subsets

showAsCaseExp :: (Show a, Show b, Eq b)
              => String -> String -> Memo a b -> String
showAsCaseExp funName varName m =
  case bindings of
    [] -> "\\" ++ varName ++ " -> " ++ funName ++ " " ++ varName ++ "\n"
    _  -> columns " " [ "\\" ++ varName ++ " -> case " ++ varName ++ " of ", caseExp ] ++ "\n"
  where caseExp = table " -> "
                . (++ [["_", funName ++ " " ++ varName]])
                . map (\(a,r) -> [show a, show r])
                $ bindings
        bindings = toList m
<commit_msg>white: fix sized instance for Int<commit_after>module FitSpec
  ( Memo (..)
  , initMemo
  , lookup
  , insert
  , toList

  , Listate (..)
  , Listate2 (..)
  , Listate3 (..)
  , runListate
  , runListate2
  , runListate3

  , lsMutateApply

  , lsMutateApply12
  , lsMutateApply22

  , lsMutateApply13
  , lsMutateApply23
  , lsMutateApply33

  , Sized (..)

  , Args (..)
  , args

  , lholds

  , report
  , reportWith
  , report2
  , report2With
  , report3
  , report3With

  , fitspec
  )
where

import Test.Check
import Control.Monad
import Data.List (transpose,sortBy)
import Data.Ord (comparing)
import Data.Maybe (isJust, isNothing, catMaybes, listToMaybe)
import Table
import qualified Data.Map as Map (insert,lookup,empty,Map,toList)
import Prelude hiding (lookup)
import Control.Applicative (liftA, liftA2)
import Data.Functor.Compose
import Utils


-- In the end, this can be merged with Listable.  It can be even generated by
-- the same combinators: in a similar way to what Uniplate does when generating
-- its different traversing functions
class Sized a where
  size :: a -> Int

instance Sized () where
  size _ = 0

instance Sized Int where
  size 0 = 0
  size n | n > 0     = n*2 - 1
         | otherwise = (abs n)*2

instance Sized Bool where
  size _ = 0

instance (Sized a) => Sized [a] where
  size = sum . map ((+1) . size)

instance (Sized a, Sized b) => Sized (a,b) where
  size (x,y) = size x + size y

instance (Sized a, Sized b, Sized c) => Sized (a,b,c) where
  size (x,y,z) = size x + size y + size z

data Memo a b = Memo (a->b) (Map.Map a (Maybe b))

initMemo :: (a->b) -> (Memo a b)
initMemo f = Memo f Map.empty

lookup :: Ord a => a -> Memo a b -> Maybe (Maybe b)
lookup x (Memo f bs) = Map.lookup x bs

insert :: Ord a => a -> Maybe b -> Memo a b -> Memo a b
insert x y (Memo f bs) = Memo f (Map.insert x y bs)

toList :: Memo a b -> [(a,b)]
toList (Memo _ bs) = map (\(x,Just y) -> (x,y))
                   . filter (\(_,m) -> isJust m)
                   $ Map.toList bs

instance (Eq b, Show a, Show b) => Show (Memo a b) where
  show m = ("Memo" ++) . show $ toList m

data Listate a b c = Listate { listates :: Memo a b
                                        -> [[ (c,Memo a b) ]] -- a listing
                             }

data Listate2 a b c d e = Listate2 { listates2 :: Memo a b
                                               -> Memo c d
                                               -> [[ (e,Memo a b,Memo c d) ]] }

data Listate3 a b c d e f g = Listate3 { listates3 :: Memo a b
                                                   -> Memo c d
                                                   -> Memo e f
                                                   -> [[ (g,Memo a b,Memo c d,Memo e f) ]] }

runListate :: (a->b)
           -> Listate a b c
           -> [[ (c,Memo a b) ]]
runListate f = (`listates` initMemo f)


runListate2 :: (a->b)
            -> (c->d)
            -> Listate2 a b c d e
            -> [[ (e,Memo a b,Memo c d) ]]
runListate2 f g = \ls -> listates2 ls (initMemo f) (initMemo g)


runListate3 :: (a->b)
            -> (c->d)
            -> (e->f)
            -> Listate3 a b c d e f g
            -> [[ (g,Memo a b,Memo c d,Memo e f) ]]
runListate3 f g h = \ls -> listates3 ls (initMemo f) (initMemo g) (initMemo h)


lsMutateApply :: (Ord a, Eq b, Sized a, Listable b)
              => a
              -> Listate a b b
lsMutateApply x = Listate $ lsMutateApply' (,) x


lsMutateApply' :: (Eq b, Ord a, Listable b, Sized a)
               => (b -> Memo a b -> t)
               -> a
               -> Memo a b
               -> [[t]]
lsMutateApply' tup x (m@(Memo f bs)) =
  case lookup x m of
    Nothing  -> [f x `tup` insert x Nothing m]
              : case errorToNothing (f x) of
                  Just (fx) -> replicate (size x) []
                            ++ lsmap (\y -> y `tup` insert x (Just y) m) (lsfilter (/= fx) listing)
                  Nothing   -> []  -- we don't mutate undefined / error values
    (Just Nothing)  -> [[ f x `tup` m ]]
    (Just (Just y)) -> [[ y `tup` m ]]


lsMutateApply12 :: (Ord a, Eq b, Sized a, Listable b)
                => a
                -> Listate2 a b c d b
lsMutateApply12 x = Listate2 $ \m m' -> lsMutateApply' (\x m'' -> (x,m'',m')) x m

lsMutateApply22 :: (Ord c, Eq d, Sized c, Listable d)
                => c
                -> Listate2 a b c d d
lsMutateApply22 x = Listate2 $ \m' m -> lsMutateApply' (\x m'' -> (x,m',m'')) x m


lsMutateApply13 :: (Ord a, Eq b, Sized a, Listable b)
                => a
                -> Listate3 a b c d e f b
lsMutateApply13 x = Listate3 $ \m n o -> lsMutateApply' (\x m' -> (x,m',n,o)) x m

lsMutateApply23 :: (Ord c, Eq d, Sized c, Listable d)
                => c
                -> Listate3 a b c d e f d
lsMutateApply23 x = Listate3 $ \m n o -> lsMutateApply' (\x n' -> (x,m,n',o)) x n

lsMutateApply33 :: (Ord e, Eq f, Sized e, Listable f)
                => e
                -> Listate3 a b c d e f f
lsMutateApply33 x = Listate3 $ \m n o -> lsMutateApply' (\x o' -> (x,m,n,o')) x o


instance Functor (Listate a b) where
  fmap = liftM

instance Applicative (Listate a b) where
  pure = return
  (<*>) = ap

instance Monad (Listate a b) where
  return x = Listate $ \m -> [[ (x,m) ]]
  ff >>= gg = Listate $ \m -> lsConcatMap (\(x,m') -> listates (gg x) m') $ listates ff m

instance Functor (Listate2 a b c d) where
  fmap = liftM

instance Applicative (Listate2 a b c d) where
  pure = return
  (<*>) = ap

instance Monad (Listate2 a b c d) where
  return x = Listate2 $ \m m' -> [[ (x,m,m') ]]
  ff >>= gg = Listate2 $ \m m' -> lsConcatMap (\(x,m'',m''') -> listates2 (gg x) m'' m''') $ listates2 ff m m'

instance Functor (Listate3 a b c d e f) where
  fmap = liftM

instance Applicative (Listate3 a b c d e f) where
  pure = return
  (<*>) = ap

instance Monad (Listate3 a b c d e f) where
  return x = Listate3 $ \m m' m'' -> [[ (x,m,m',m'') ]]
  ff >>= gg = Listate3
            $ \m m' m'' -> lsConcatMap (\(x,m''',m'''',m''''') -> listates3 (gg x) m''' m'''' m''''')
                                     $ listates3 ff m m' m''

ltest :: Monad m
      => (c -> m Bool)
      -> [c]
      -> m Bool
ltest prop xs = foldr (liftM2 (&&)) (return True) (map prop xs)

lholds :: (Monad m, Listable c)
       => Int
       -> (c -> m Bool)
       -> m Bool
lholds n prop = ltest prop (take n list)


-- | Extra arguments for reporting functions
data Args = Args
          { functionNames :: [String]
          , variableNames :: [String]
          , limitResults :: Maybe Int -- maximum number of results
          }

-- | Default arguments
args = Args { functionNames = repeat "function"
            , variableNames = repeat "x"
            , limitResults = Nothing    -- show everything
            }

report :: (Show a, Show b, Eq b) => Int -> [[([Bool], Memo a b)]] -> IO ()
report nf rss = reportWith args nf rss

reportWith :: (Show a, Show b, Eq b) => Args -> Int -> [[([Bool], Memo a b)]] -> IO ()
reportWith args nf rss = putStrLn
                       . table "   "
                       . map showResult
                       . maybe id take (limitResults args)
                       $ fitspec fst snd nf rss
  where showResult (x,y,mm) = [ show x, show y, showM mm ]
        showM Nothing  = ""
        showM (Just m) = showAsCaseExp (functionNames args !! 0) (variableNames args !! 0) m

report2 :: (Show a, Show b, Eq b, Show c, Show d, Eq d)
        => Int -> [[([Bool], Memo a b, Memo c d)]] -> IO ()
report2 nf rss = report2With args nf rss

report2With :: (Show a, Show b, Eq b, Show c, Show d, Eq d)
            => Args -> Int -> [[([Bool], Memo a b, Memo c d)]] -> IO ()
report2With args nf rss = putStrLn
                        . table "   "
                        . map showResult
                        . maybe id take (limitResults args)
                        $ fitspec (\(x,_,_)->x) (\(_,m,n)->(m,n)) nf rss
  where showResult (x,y,mm) = [ show x, show y, showM mm ]
        showM Nothing  = ""
        showM (Just (m,n)) = showAsCaseExp (functionNames args !! 0) (variableNames args !! 0) m
                          ++ showAsCaseExp (functionNames args !! 1) (variableNames args !! 1) n

report3 :: (Show a, Show b, Eq b, Show c, Show d, Eq d, Show e, Show f, Eq f)
        => Int -> [[([Bool], Memo a b, Memo c d, Memo e f)]] -> IO ()
report3 nf rss = report3With args nf rss

report3With :: (Show a, Show b, Eq b, Show c, Show d, Eq d, Show e, Show f, Eq f)
            => Args -> Int -> [[([Bool], Memo a b, Memo c d, Memo e f)]] -> IO ()
report3With args nf rss = putStrLn
                        . table "   "
                        . map showResult
                        . maybe id take (limitResults args)
                        $ fitspec (\(x,_,_,_)->x) (\(_,m,n,o)->(m,n,o)) nf rss
  where showResult (x,y,mm) = [ show x, show y, showM mm ]
        showM Nothing  = ""
        showM (Just (m,n,o)) = showAsCaseExp (functionNames args !! 0) (variableNames args !! 0) m
                            ++ showAsCaseExp (functionNames args !! 1) (variableNames args !! 1) n
                            ++ showAsCaseExp (functionNames args !! 2) (variableNames args !! 2) o

fitspec :: (a -> [Bool]) -> (a -> b)
        -> Int -> [[a]] -> [(Int, [Int], Maybe b)]
fitspec getrs getsm n rss@((p:_):_) =
      filterU relevant
    . sortBy (comparing $ \(s,ids,_) -> (s, length ids))
    . tag
    . fitspec' getrs getsm n
    $ rss
  where tag = zipWith (\ids (ns,sms) -> (ns,ids,sms)) (subsets [1..length (getrs p)])
        (n,is,_) `relevant` (n',is',_) = not (n == n' && is `contained` is')
fitspec _ _ _ _ = error "fitspec: something went wrong, no results to analyse?"

fitspec' :: (a -> [Bool]) -> (a -> b)
         -> Int -> [[a]] -> [(Int, Maybe b)]
fitspec' getrs getsm n =
      map (\fs -> (count isJust fs, listToMaybe $ catMaybes fs))
    . transpose
    . map (\p -> map (\b -> if b then Just (getsm p) else Nothing) (compositions (getrs p)))
    . take n . tail . concat
  where count f = length . filter f


-- | 'compositions' @bs@ returns all compositions formed by taking values of @bs@
compositions :: [Bool] -> [Bool]
compositions = map and . subsets

showAsCaseExp :: (Show a, Show b, Eq b)
              => String -> String -> Memo a b -> String
showAsCaseExp funName varName m =
  case bindings of
    [] -> "\\" ++ varName ++ " -> " ++ funName ++ " " ++ varName ++ "\n"
    _  -> columns " " [ "\\" ++ varName ++ " -> case " ++ varName ++ " of ", caseExp ] ++ "\n"
  where caseExp = table " -> "
                . (++ [["_", funName ++ " " ++ varName]])
                . map (\(a,r) -> [show a, show r])
                $ bindings
        bindings = toList m

{'fraction_non_alphanumeric': 0.14925302794665157, 'fraction_numerical': 0.010038173335218436, 'mean_word_length': 1.7844114945545204, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 311, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 28, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '868765', 'n_tokens_mistral': 7613, 'n_tokens_neox': 7319, 'n_words': 2658}