-- | Miscellaneous utility functions.
{-# OPTIONS_HADDOCK hide #-}
{-# LANGUAGE CPP #-}
module QuickSpec.Internal.Utils where

import Control.Arrow((&&&))
import Control.Exception
import Control.Spoon
import Data.List(groupBy, sortBy)
#if !MIN_VERSION_base(4,8,0)
import Data.Monoid
#endif
import Data.Ord(comparing)
import System.IO
import qualified Control.Category as Category
import qualified Data.Map.Strict as Map
import Data.Map(Map)
import Language.Haskell.TH.Syntax
import Data.Lens.Light
import Twee.Base hiding (lookup)
import Control.Monad.Trans.State.Strict
import Control.Monad

(#) :: Category.Category cat => cat b c -> cat a b -> cat a c
# :: forall (cat :: * -> * -> *) b c a.
Category cat =>
cat b c -> cat a b -> cat a c
(#) = forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
(Category..)

key :: Ord a => a -> Lens (Map a b) (Maybe b)
key :: forall a b. Ord a => a -> Lens (Map a b) (Maybe b)
key a
x = forall a b. (a -> b) -> (b -> a -> a) -> Lens a b
lens (forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
x) (\Maybe b
my Map a b
m -> forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (forall a b. a -> b -> a
const Maybe b
my) a
x Map a b
m)

keyDefault :: Ord a => a -> b -> Lens (Map a b) b
keyDefault :: forall a b. Ord a => a -> b -> Lens (Map a b) b
keyDefault a
x b
y = forall a b. (a -> b) -> (b -> a -> a) -> Lens a b
lens (forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault b
y a
x) (\b
y Map a b
m -> forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
x b
y Map a b
m)

reading :: (a -> Lens a b) -> Lens a b
reading :: forall a b. (a -> Lens a b) -> Lens a b
reading a -> Lens a b
f = forall a b. (a -> b) -> (b -> a -> a) -> Lens a b
lens (\a
x -> forall a b. Lens a b -> a -> b
getL (a -> Lens a b
f a
x) a
x) (\b
y a
x -> forall a b. Lens a b -> b -> a -> a
setL (a -> Lens a b
f a
x) b
y a
x)

fstLens :: Lens (a, b) a
fstLens :: forall a b. Lens (a, b) a
fstLens = forall a b. (a -> b) -> (b -> a -> a) -> Lens a b
lens forall a b. (a, b) -> a
fst (\a
x (a
_, b
y) -> (a
x, b
y))

sndLens :: Lens (a, b) b
sndLens :: forall a b. Lens (a, b) b
sndLens = forall a b. (a -> b) -> (b -> a -> a) -> Lens a b
lens forall a b. (a, b) -> b
snd (\b
y (a
x, b
_) -> (a
x, b
y))

makeLensAs :: Name -> [(String, String)] -> Q [Dec]
makeLensAs :: Name -> [(String, String)] -> Q [Dec]
makeLensAs Name
ty [(String, String)]
names =
  Name -> (String -> Maybe String) -> Q [Dec]
nameMakeLens Name
ty (\String
x -> forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
x [(String, String)]
names)

repeatM :: Monad m => m a -> m [a]
repeatM :: forall (m :: * -> *) a. Monad m => m a -> m [a]
repeatM = forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> [a]
repeat

partitionBy :: Ord b => (a -> b) -> [a] -> [[a]]
partitionBy :: forall b a. Ord b => (a -> b) -> [a] -> [[a]]
partitionBy a -> b
value =
  forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
  forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (\(a, b)
x (a, b)
y -> forall a b. (a, b) -> b
snd (a, b)
x forall a. Eq a => a -> a -> Bool
== forall a b. (a, b) -> b
snd (a, b)
y) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
  forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing forall a b. (a, b) -> b
snd) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
  forall a b. (a -> b) -> [a] -> [b]
map (forall a. a -> a
id forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& a -> b
value)

collate :: Ord a => ([b] -> c) -> [(a, b)] -> [(a, c)]
collate :: forall a b c. Ord a => ([b] -> c) -> [(a, b)] -> [(a, c)]
collate [b] -> c
f = forall a b. (a -> b) -> [a] -> [b]
map forall {a}. [(a, b)] -> (a, c)
g forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a. Ord b => (a -> b) -> [a] -> [[a]]
partitionBy forall a b. (a, b) -> a
fst
  where
    g :: [(a, b)] -> (a, c)
g [(a, b)]
xs = (forall a b. (a, b) -> a
fst (forall a. [a] -> a
head [(a, b)]
xs), [b] -> c
f (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(a, b)]
xs))

isSorted :: Ord a => [a] -> Bool
isSorted :: forall a. Ord a => [a] -> Bool
isSorted [a]
xs = forall (t :: * -> *). Foldable t => t Bool -> Bool
and (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall a. Ord a => a -> a -> Bool
(<=) [a]
xs (forall a. [a] -> [a]
tail [a]
xs))

isSortedBy :: Ord b => (a -> b) -> [a] -> Bool
isSortedBy :: forall b a. Ord b => (a -> b) -> [a] -> Bool
isSortedBy a -> b
f [a]
xs = forall a. Ord a => [a] -> Bool
isSorted (forall a b. (a -> b) -> [a] -> [b]
map a -> b
f [a]
xs)

usort :: Ord a => [a] -> [a]
usort :: forall a. Ord a => [a] -> [a]
usort = forall a. (a -> a -> Ordering) -> [a] -> [a]
usortBy forall a. Ord a => a -> a -> Ordering
compare

usortBy :: (a -> a -> Ordering) -> [a] -> [a]
usortBy :: forall a. (a -> a -> Ordering) -> [a] -> [a]
usortBy a -> a -> Ordering
f = forall a b. (a -> b) -> [a] -> [b]
map forall a. [a] -> a
head forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (\a
x a
y -> a -> a -> Ordering
f a
x a
y forall a. Eq a => a -> a -> Bool
== Ordering
EQ) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy a -> a -> Ordering
f

sortBy' :: Ord b => (a -> b) -> [a] -> [a]
sortBy' :: forall b a. Ord b => (a -> b) -> [a] -> [a]
sortBy' a -> b
f = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a -> b
f a
x, a
x))

usortBy' :: Ord b => (a -> b) -> [a] -> [a]
usortBy' :: forall b a. Ord b => (a -> b) -> [a] -> [a]
usortBy' a -> b
f = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Ordering) -> [a] -> [a]
usortBy (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (\a
x -> (a -> b
f a
x, a
x))

orElse :: Ordering -> Ordering -> Ordering
Ordering
EQ orElse :: Ordering -> Ordering -> Ordering
`orElse` Ordering
x = Ordering
x
Ordering
x  `orElse` Ordering
_ = Ordering
x

unbuffered :: IO a -> IO a
unbuffered :: forall a. IO a -> IO a
unbuffered IO a
x = do
  BufferMode
buf <- Handle -> IO BufferMode
hGetBuffering Handle
stdout
  forall a b c. IO a -> IO b -> IO c -> IO c
bracket_
    (Handle -> BufferMode -> IO ()
hSetBuffering Handle
stdout BufferMode
NoBuffering)
    (Handle -> BufferMode -> IO ()
hSetBuffering Handle
stdout BufferMode
buf)
    IO a
x

spoony :: Eq a => a -> Maybe a
spoony :: forall a. Eq a => a -> Maybe a
spoony a
x = forall a. a -> Maybe a
teaspoon ((a
x forall a. Eq a => a -> a -> Bool
== a
x) seq :: forall a b. a -> b -> b
`seq` a
x)

labelM :: Monad m => (a -> m b) -> [a] -> m [(a, b)]
labelM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> [a] -> m [(a, b)]
labelM a -> m b
f = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\a
x -> do { b
y <- a -> m b
f a
x; forall (m :: * -> *) a. Monad m => a -> m a
return (a
x, b
y) })

#if __GLASGOW_HASKELL__ < 710
isSubsequenceOf :: Ord a => [a] -> [a] -> Bool
[] `isSubsequenceOf` ys = True
(x:xs) `isSubsequenceOf` [] = False
(x:xs) `isSubsequenceOf` (y:ys)
  | x == y = xs `isSubsequenceOf` ys
  | otherwise = (x:xs) `isSubsequenceOf` ys
#endif

appendAt :: Int -> [a] -> [[a]] -> [[a]]
appendAt :: forall a. Int -> [a] -> [[a]] -> [[a]]
appendAt Int
n [a]
xs [] = forall a. Int -> [a] -> [[a]] -> [[a]]
appendAt Int
n [a]
xs [[]]
appendAt Int
0 [a]
xs ([a]
ys:[[a]]
yss) = ([a]
ys forall a. [a] -> [a] -> [a]
++ [a]
xs)forall a. a -> [a] -> [a]
:[[a]]
yss
appendAt Int
n [a]
xs ([a]
ys:[[a]]
yss) = [a]
ysforall a. a -> [a] -> [a]
:forall a. Int -> [a] -> [[a]] -> [[a]]
appendAt (Int
nforall a. Num a => a -> a -> a
-Int
1) [a]
xs [[a]]
yss

-- Should be in Twee.Base.
antiunify :: Ord f => Term f -> Term f -> Term f
antiunify :: forall f. Ord f => Term f -> Term f -> Term f
antiunify Term f
t Term f
u =
  forall a. Build a => a -> Term (BuildFun a)
build forall a b. (a -> b) -> a -> b
$ forall s a. State s a -> s -> a
evalState (forall {m :: * -> *} {f}.
Monad m =>
Term f
-> Term f -> StateT (Var, Map (Term f, Term f) Var) m (Builder f)
loop Term f
t Term f
u) (forall a. Enum a => a -> a
succ (forall a b. (a, b) -> b
snd (forall f. Term f -> (Var, Var)
bound Term f
t) forall a. Ord a => a -> a -> a
`max` forall a b. (a, b) -> b
snd (forall f. Term f -> (Var, Var)
bound Term f
u)), forall k a. Map k a
Map.empty)
  where
    loop :: Term f
-> Term f -> StateT (Var, Map (Term f, Term f) Var) m (Builder f)
loop (App Fun f
f TermList f
ts) (App Fun f
g TermList f
us)
      | Fun f
f forall a. Eq a => a -> a -> Bool
== Fun f
g =
        forall a. Build a => Fun (BuildFun a) -> a -> Builder (BuildFun a)
app Fun f
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM Term f
-> Term f -> StateT (Var, Map (Term f, Term f) Var) m (Builder f)
loop (forall f. TermList f -> [Term f]
unpack TermList f
ts) (forall f. TermList f -> [Term f]
unpack TermList f
us)
    loop (Var Var
x) (Var Var
y)
      | Var
x forall a. Eq a => a -> a -> Bool
== Var
y =
        forall (m :: * -> *) a. Monad m => a -> m a
return (forall f. Var -> Builder f
var Var
x)
    loop Term f
t Term f
u = do
      (Var
next, Map (Term f, Term f) Var
m) <- forall (m :: * -> *) s. Monad m => StateT s m s
get
      case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Term f
t, Term f
u) Map (Term f, Term f) Var
m of
        Just Var
v -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall f. Var -> Builder f
var Var
v)
        Maybe Var
Nothing -> do
          forall (m :: * -> *) s. Monad m => s -> StateT s m ()
put (forall a. Enum a => a -> a
succ Var
next, forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (Term f
t, Term f
u) Var
next Map (Term f, Term f) Var
m)
          forall (m :: * -> *) a. Monad m => a -> m a
return (forall f. Var -> Builder f
var Var
next)

{-# INLINE fixpoint #-}
fixpoint :: Eq a => (a -> a) -> a -> a
fixpoint :: forall a. Eq a => (a -> a) -> a -> a
fixpoint a -> a
f a
x = a -> a
fxp a
x
  where
    fxp :: a -> a
fxp a
x
      | a
x forall a. Eq a => a -> a -> Bool
== a
y = a
x
      | Bool
otherwise = a -> a
fxp a
y
      where
        y :: a
y = a -> a
f a
x