{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeApplications #-}

module Kanren.Example.List (example) where

import Data.Function ((&))
import Data.Maybe (fromJust)

import Kanren.Core
import Kanren.Goal
import Kanren.LogicalBase
import Kanren.Match

listo :: (Logical a) => Term [a] -> Goal ()
listo :: forall a. Logical a => Term [a] -> Goal ()
listo =
  Term [a] -> Goal ()
forall a x. Term a -> Goal x
matche
    (Term [a] -> Goal ())
-> ((Term [a] -> Goal ()) -> Term [a] -> Goal ())
-> Term [a]
-> Goal ()
forall a b. a -> (a -> b) -> b
& Prism' (Logic [a]) ()
-> (() -> Goal ()) -> (Term [a] -> Goal ()) -> Term [a] -> Goal ()
forall a v x.
(Logical a, Fresh v) =>
Prism' (Logic a) v
-> (v -> Goal x) -> (Term a -> Goal x) -> Term a -> Goal x
on p () (f ()) -> p (Logic [a]) (f (Logic [a]))
p () (f ()) -> p (LogicList a) (f (LogicList a))
forall a (p :: * -> * -> *) (f :: * -> *).
(Choice p, Applicative f) =>
p () (f ()) -> p (LogicList a) (f (LogicList a))
Prism' (Logic [a]) ()
_LogicNil () -> Goal ()
forall a. a -> Goal a
forall (m :: * -> *) a. Monad m => a -> m a
return
    (Term [a] -> Goal ())
-> ((Term [a] -> Goal ()) -> Term [a] -> Goal ())
-> Term [a]
-> Goal ()
forall a b. a -> (a -> b) -> b
& Prism' (Logic [a]) (Term a, Term [a])
-> ((Term a, Term [a]) -> Goal ())
-> (Term [a] -> Goal ())
-> Term [a]
-> Goal ()
forall a v x.
(Logical a, Fresh v) =>
Prism' (Logic a) v
-> (v -> Goal x) -> (Term a -> Goal x) -> Term a -> Goal x
on p (Term a, Term [a]) (f (Term a, Term [a]))
-> p (Logic [a]) (f (Logic [a]))
p (Term a, Term [a]) (f (Term a, Term [a]))
-> p (LogicList a) (f (LogicList a))
forall a1 a2 (p :: * -> * -> *) (f :: * -> *).
(Choice p, Applicative f) =>
p (Term a1, Term [a1]) (f (Term a2, Term [a2]))
-> p (LogicList a1) (f (LogicList a2))
Prism' (Logic [a]) (Term a, Term [a])
_LogicCons (\(Term a
_, Term [a]
xs) -> Term [a] -> Goal ()
forall a. Logical a => Term [a] -> Goal ()
listo Term [a]
xs)

{- FOURMOLU_DISABLE -}
appendo :: (Logical a) => Term [a] -> Term [a] -> Term [a] -> Goal ()
appendo :: forall a. Logical a => Term [a] -> Term [a] -> Term [a] -> Goal ()
appendo Term [a]
xs Term [a]
ys Term [a]
zs = Term [a]
xs Term [a] -> (Term [a] -> Goal ()) -> Goal ()
forall a b. a -> (a -> b) -> b
& (Term [a] -> Goal ()
forall a x. Term a -> Goal x
matche
  (Term [a] -> Goal ())
-> ((Term [a] -> Goal ()) -> Term [a] -> Goal ())
-> Term [a]
-> Goal ()
forall a b. a -> (a -> b) -> b
& Prism' (Logic [a]) ()
-> (() -> Goal ()) -> (Term [a] -> Goal ()) -> Term [a] -> Goal ()
forall a v x.
(Logical a, Fresh v) =>
Prism' (Logic a) v
-> (v -> Goal x) -> (Term a -> Goal x) -> Term a -> Goal x
on p () (f ()) -> p (Logic [a]) (f (Logic [a]))
p () (f ()) -> p (LogicList a) (f (LogicList a))
forall a (p :: * -> * -> *) (f :: * -> *).
(Choice p, Applicative f) =>
p () (f ()) -> p (LogicList a) (f (LogicList a))
Prism' (Logic [a]) ()
_LogicNil (\() -> Term [a]
ys Term [a] -> Term [a] -> Goal ()
forall a. Logical a => Term a -> Term a -> Goal ()
=== Term [a]
zs)
  (Term [a] -> Goal ())
-> ((Term [a] -> Goal ()) -> Term [a] -> Goal ())
-> Term [a]
-> Goal ()
forall a b. a -> (a -> b) -> b
& Prism' (Logic [a]) (Term a, Term [a])
-> ((Term a, Term [a]) -> Goal ())
-> (Term [a] -> Goal ())
-> Term [a]
-> Goal ()
forall a v x.
(Logical a, Fresh v) =>
Prism' (Logic a) v
-> (v -> Goal x) -> (Term a -> Goal x) -> Term a -> Goal x
on p (Term a, Term [a]) (f (Term a, Term [a]))
-> p (Logic [a]) (f (Logic [a]))
p (Term a, Term [a]) (f (Term a, Term [a]))
-> p (LogicList a) (f (LogicList a))
forall a1 a2 (p :: * -> * -> *) (f :: * -> *).
(Choice p, Applicative f) =>
p (Term a1, Term [a1]) (f (Term a2, Term [a2]))
-> p (LogicList a1) (f (LogicList a2))
Prism' (Logic [a]) (Term a, Term [a])
_LogicCons (\(Term a
x, Term [a]
xs') -> do
      Term [a]
zs' <- Goal (Term [a])
forall v. Fresh v => Goal v
fresh
      Term [a]
zs Term [a] -> Term [a] -> Goal ()
forall a. Logical a => Term a -> Term a -> Goal ()
=== Logic [a] -> Term [a]
forall a. Logic a -> Term a
Value (Term a -> Term [a] -> LogicList a
forall a. Term a -> Term [a] -> LogicList a
LogicCons Term a
x Term [a]
zs')
      Term [a] -> Term [a] -> Term [a] -> Goal ()
forall a. Logical a => Term [a] -> Term [a] -> Term [a] -> Goal ()
appendo Term [a]
xs' Term [a]
ys Term [a]
zs'))
{- FOURMOLU_ENABLE -}

partitions :: (Logical a) => [a] -> [([a], [a])]
partitions :: forall a. Logical a => [a] -> [([a], [a])]
partitions [a]
xs = (Term [a], Term [a]) -> ([a], [a])
reifyBoth ((Term [a], Term [a]) -> ([a], [a]))
-> [(Term [a], Term [a])] -> [([a], [a])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Term [a], Term [a])]
partitioned
 where
  partitioned :: [(Term [a], Term [a])]
partitioned = ((Term [a], Term [a]) -> Goal ()) -> [(Term [a], Term [a])]
forall v. Fresh v => (v -> Goal ()) -> [v]
run (((Term [a], Term [a]) -> Goal ()) -> [(Term [a], Term [a])])
-> ((Term [a], Term [a]) -> Goal ()) -> [(Term [a], Term [a])]
forall a b. (a -> b) -> a -> b
$ \(Term [a]
left, Term [a]
right) -> do
    Term [a] -> Term [a] -> Term [a] -> Goal ()
forall a. Logical a => Term [a] -> Term [a] -> Term [a] -> Goal ()
appendo Term [a]
left Term [a]
right ([a] -> Term [a]
forall a. Logical a => a -> Term a
inject' [a]
xs)

  reifyBoth :: (Term [a], Term [a]) -> ([a], [a])
reifyBoth (Term [a]
a, Term [a]
b) = (Term [a] -> [a]
reify Term [a]
a, Term [a] -> [a]
reify Term [a]
b)
  reify :: Term [a] -> [a]
reify = Maybe [a] -> [a]
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe [a] -> [a]) -> (Term [a] -> Maybe [a]) -> Term [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term [a] -> Maybe [a]
forall a. Logical a => Term a -> Maybe a
extract'

showLogicList :: (Show (Logic a)) => Term [a] -> String
showLogicList :: forall a. Show (Logic a) => Term [a] -> String
showLogicList Term [a]
list = String
"[" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Term [a] -> String
forall a. Show (Logic a) => Term [a] -> String
go Term [a]
list String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"]"
 where
  go :: Term [a] -> String
go (Var VarId [a]
_) = String
"..?"
  go (Value Logic [a]
LogicList a
LogicNil) = String
""
  go (Value (LogicCons Term a
x Term [a]
xs)) = Term a -> String
forall a. Show a => a -> String
show Term a
x String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
sep String -> String -> String
forall a. [a] -> [a] -> [a]
++ Term [a] -> String
go Term [a]
xs
   where
    sep :: String
sep = case Term [a]
xs of
      Value Logic [a]
LogicList a
LogicNil -> String
""
      Term [a]
_ -> String
", "

example :: IO ()
example :: IO ()
example = do
  String -> IO ()
putStrLn String
"listo:"
  (Term [Int] -> IO ()) -> [Term [Int]] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (String -> IO ()
putStrLn (String -> IO ()) -> (Term [Int] -> String) -> Term [Int] -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term [Int] -> String
forall a. Show (Logic a) => Term [a] -> String
showLogicList) ([Term [Int]] -> IO ()) -> [Term [Int]] -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> [Term [Int]] -> [Term [Int]]
forall a. Int -> [a] -> [a]
take Int
5 ((Term [Int] -> Goal ()) -> [Term [Int]]
forall v. Fresh v => (v -> Goal ()) -> [v]
run (forall a. Logical a => Term [a] -> Goal ()
listo @Int))

  String -> IO ()
putStrLn String
"\npartitions [1, 2, 3]:"
  (([Int], [Int]) -> IO ()) -> [([Int], [Int])] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ([Int], [Int]) -> IO ()
forall a. Show a => a -> IO ()
print ([([Int], [Int])] -> IO ()) -> [([Int], [Int])] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Int] -> [([Int], [Int])]
forall a. Logical a => [a] -> [([a], [a])]
partitions [Int
1 :: Int, Int
2, Int
3]