{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}

module Data.PreNatMap(
  PreNatMap(),

  -- * initialize
  empty,
  
  -- * query
  isFull,
  fullMatch, match,
  lookup, lookupWith, lookupShape,

  -- * update
  refine, refineShape,
  refine', refineShape',

  -- * conversion
  toEntries, fromEntries, make,
  toNatMap, fromNatMap, toShapeMap, fromShapeMap,
) where

import Prelude hiding (lookup)

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import qualified Data.IntMap.Strict as IntMap
import Data.Foldable (Foldable(..))

import Data.FunctorShape
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector as V
import Control.Monad (guard)
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Functor.Classes (showsUnaryWith)
import Data.NatMap (NatMap)
import qualified Data.NatMap as NM

import TraversableUtil (indices, zipMatch)
import qualified Data.Set as Set

-- | @PreNatMap f g@ represents a partially known
--   natural transformation @f ~> g@.
-- 
-- ==== Examples
-- 
-- @reverse@ is a natural transformation from list to list.
--
-- > reverse :: [a] -> [a]
-- 
-- Using 'refine' function, you can accumulate knowledge about how
-- @reverse@ behaves to @PreNatMap [] []@ through examples.
-- 
-- >>> p1 = refine "foo" "oof" empty
-- >>> p1
-- Just (make [([0,1,1],[1,1,0])])
--
-- You can add another example to strengthen the knowledge.
-- (The monadic bind @'>>='@ is used here because 'refine'
--  returns the result wrapped in @Maybe@.)
--
-- >>> p2 = p1 >>= refine "dad" "dad"
-- >>> p2
-- Just (make [([0,1,2],[2,1,0])])
--
-- >>> p3 = p2 >>= refine "aabb" "bbaa"
-- >>> p3
-- Just (make [([0,1,2],[2,1,0]),([0,0,1,1],[1,1,0,0])])
--
-- If you add an example contradicting to existing examples,
-- 'refine' can fail with @Nothing@.
-- 
-- >>> p3 >>= refine "cccd" "cddd"
-- Nothing
-- 
-- You can also query against the learned knowledge.
-- 'match' function takes an input and a 'PreNatMap', and
-- if the learned knowledge determine the output uniquely,
-- returns that output.
--
-- >>> p3 >>= match "ABC"
-- Just "CBA"
-- >>> p3 >>= match "XXYY"
-- Just "YYXX"
-- 
-- If the output is not unique, 'match' can fail,
-- even when the "shape" of the output is known.
-- In the following example, it is known that the output for a 4-lengthed list
-- is a 4-lengthed list again, but 'match' fails because the differences
-- in the list elements make the output ambiguous.
--
-- >>> p3 >>= match "XYZW"
-- Nothing
-- 
-- There is another query 'lookup' which always return
-- something for "at least the shape is known" case,
-- by combining ambiguous inputs using 'Semigroup' operation.
-- 
-- >>> p3 >>= lookup ["X", "Y", "Z", "W"] 
-- Just ["WZ","WZ","YX","YX"]
newtype PreNatMap f g = PreNatMap (Map (Shape f) (PosData g))

deriving instance (Eq (f Ignored), Eq (g Int)) => Eq (PreNatMap f g)
deriving instance (Ord (f Ignored), Ord (g Int)) => Ord (PreNatMap f g)

instance (Show (f Int), Show (g Int), Traversable f, Functor g) => Show (PreNatMap f g) where
  showsPrec :: Int -> PreNatMap f g -> ShowS
showsPrec Int
p PreNatMap f g
pnm = (Int -> [(f Int, g Int)] -> ShowS)
-> String -> Int -> [(f Int, g Int)] -> ShowS
forall a. (Int -> a -> ShowS) -> String -> Int -> a -> ShowS
showsUnaryWith Int -> [(f Int, g Int)] -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec String
"make" Int
p (PreNatMap f g -> [(f Int, g Int)]
forall (f :: * -> *) (g :: * -> *).
(Traversable f, Functor g) =>
PreNatMap f g -> [(f Int, g Int)]
toEntries PreNatMap f g
pnm)

data PosData g = PosData
  !(UV.Vector Int)
    -- ^ Vector @v@ representing LHS pattern variables.
    --
    -- Invariant (canonical reindexing):
    -- scanning @v@ left-to-right, when a value appears for the first time,
    -- it must be exactly the next fresh id (0,1,2,...).
    --
    -- More formally: let @firstOcc x@ be the least index @i@ with @v!i = x@.
    -- Then @v!(firstOcc x) == #{ y | firstOcc y < firstOcc x }@.
    --
    -- Consequences (under this invariant):
    --
    -- * If @v@ is nonempty then @v!0 == 0@.
    -- * For all indices @i@, @0 <= v!i <= i@.
    -- * If @v@ has @k@ distinct values then @maximum v == k-1@.
    -- * All elements of @v@ are pairwise distinct iff @v == [0..n-1]@
    --   (where @n = length v@). Equivalently, @n==0 || v!(n-1) == n-1@;
    --   see 'isCompleteLHS'.

  !(g Int)
    -- ^ RHS variable expression @gx@.
    --
    -- Invariant: every variable occurring in @gx@ must also occur in @v@.

deriving instance (Eq (g Int)) => Eq (PosData g)
deriving instance (Ord (g Int)) => Ord (PosData g)
deriving instance (Show (g Int)) => Show (PosData g)

-- * Entry

-- | Empty 'PreNatMap' with no knowledge.
empty :: PreNatMap f g
empty :: forall (f :: * -> *) (g :: * -> *). PreNatMap f g
empty = Map (Shape f) (PosData g) -> PreNatMap f g
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (PosData g) -> PreNatMap f g
PreNatMap Map (Shape f) (PosData g)
forall k a. Map k a
Map.empty

-- | Extract @PreNatMap@ as a list of examples.
toEntries :: (Traversable f, Functor g) => PreNatMap f g -> [(f Int, g Int)]
toEntries :: forall (f :: * -> *) (g :: * -> *).
(Traversable f, Functor g) =>
PreNatMap f g -> [(f Int, g Int)]
toEntries (PreNatMap Map (Shape f) (PosData g)
pnm) = (Shape f, PosData g) -> (f Int, g Int)
forall {f :: * -> *} {g :: * -> *}.
Traversable f =>
(Shape f, PosData g) -> (f Int, g Int)
preEntry ((Shape f, PosData g) -> (f Int, g Int))
-> [(Shape f, PosData g)] -> [(f Int, g Int)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map (Shape f) (PosData g) -> [(Shape f, PosData g)]
forall k a. Map k a -> [(k, a)]
Map.toList Map (Shape f) (PosData g)
pnm
  where
    preEntry :: (Shape f, PosData g) -> (f Int, g Int)
preEntry (Shape f a
f, PosData Vector Int
lhs g Int
rhs) = (f Int
fn, g Int
rhs)
      where
        fi :: f Int
fi = f a -> f Int
forall (f :: * -> *) x. Traversable f => f x -> f Int
indices f a
f
        fn :: f Int
fn = (Vector Int
lhs Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
UV.!) (Int -> Int) -> f Int -> f Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f Int
fi

-- | Rebuild from a list of examples
fromEntries :: (Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g, Ord a)
  => [(f a, g a)] -> Maybe (PreNatMap f g)
fromEntries :: forall (f :: * -> *) (g :: * -> *) a.
(Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g,
 Ord a) =>
[(f a, g a)] -> Maybe (PreNatMap f g)
fromEntries = (Maybe (PreNatMap f g) -> (f a, g a) -> Maybe (PreNatMap f g))
-> Maybe (PreNatMap f g) -> [(f a, g a)] -> Maybe (PreNatMap f g)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Maybe (PreNatMap f g)
mm (f a
fa, g a
ga) -> Maybe (PreNatMap f g)
mm Maybe (PreNatMap f g)
-> (PreNatMap f g -> Maybe (PreNatMap f g))
-> Maybe (PreNatMap f g)
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ !PreNatMap f g
m -> f a -> g a -> PreNatMap f g -> Maybe (PreNatMap f g)
forall a (f :: * -> *) (g :: * -> *).
(Ord a, Ord (f Ignored), Eq (g Ignored), Foldable f,
 Traversable g) =>
f a -> g a -> PreNatMap f g -> Maybe (PreNatMap f g)
refine f a
fa g a
ga PreNatMap f g
m) (PreNatMap f g -> Maybe (PreNatMap f g)
forall a. a -> Maybe a
Just PreNatMap f g
forall (f :: * -> *) (g :: * -> *). PreNatMap f g
empty)

-- | 'fromEntries' but raises 'error' on 'Nothing' case.
make :: (Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g, Ord a)
  => [(f a, g a)] -> PreNatMap f g
make :: forall (f :: * -> *) (g :: * -> *) a.
(Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g,
 Ord a) =>
[(f a, g a)] -> PreNatMap f g
make = PreNatMap f g -> Maybe (PreNatMap f g) -> PreNatMap f g
forall a. a -> Maybe a -> a
fromMaybe (String -> PreNatMap f g
forall a. HasCallStack => String -> a
error String
"make: inconsistent entries") (Maybe (PreNatMap f g) -> PreNatMap f g)
-> ([(f a, g a)] -> Maybe (PreNatMap f g))
-> [(f a, g a)]
-> PreNatMap f g
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(f a, g a)] -> Maybe (PreNatMap f g)
forall (f :: * -> *) (g :: * -> *) a.
(Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g,
 Ord a) =>
[(f a, g a)] -> Maybe (PreNatMap f g)
fromEntries

-- | Convert to 'NatMap' by discarding non-full entries
toNatMap :: (Ord (f Ignored), Traversable f, Functor g) => PreNatMap f g -> NatMap f g
toNatMap :: forall (f :: * -> *) (g :: * -> *).
(Ord (f Ignored), Traversable f, Functor g) =>
PreNatMap f g -> NatMap f g
toNatMap (PreNatMap Map (Shape f) (PosData g)
pnm) = [Entry f g] -> NatMap f g
forall (f :: * -> *) (g :: * -> *).
WeakOrd f =>
[Entry f g] -> NatMap f g
NM.fromEntries ([Entry f g] -> NatMap f g) -> [Entry f g] -> NatMap f g
forall a b. (a -> b) -> a -> b
$ ((Shape f, PosData g) -> Maybe (Entry f g))
-> [(Shape f, PosData g)] -> [Entry f g]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Shape f, PosData g) -> Maybe (Entry f g)
forall {f :: * -> *} {g :: * -> *}.
(Traversable f, Functor g) =>
(Shape f, PosData g) -> Maybe (Entry f g)
fullToEntry ([(Shape f, PosData g)] -> [Entry f g])
-> [(Shape f, PosData g)] -> [Entry f g]
forall a b. (a -> b) -> a -> b
$ Map (Shape f) (PosData g) -> [(Shape f, PosData g)]
forall k a. Map k a -> [(k, a)]
Map.toList Map (Shape f) (PosData g)
pnm
  where
    fullToEntry :: (Shape f, PosData g) -> Maybe (Entry f g)
fullToEntry (Shape f a
f, pd :: PosData g
pd@(PosData Vector Int
_ g Int
rhs))
      | PosData g -> Bool
forall (g :: * -> *). PosData g -> Bool
isCompleteLHS PosData g
pd = Entry f g -> Maybe (Entry f g)
forall a. a -> Maybe a
Just (Entry f g -> Maybe (Entry f g)) -> Entry f g -> Maybe (Entry f g)
forall a b. (a -> b) -> a -> b
$ f Int -> g Int -> Entry f g
forall (f :: * -> *) (g :: * -> *) k.
(Traversable f, Functor g, Ord k) =>
f k -> g k -> Entry f g
NM.unsafeMakeEntry (f a -> f Int
forall (f :: * -> *) x. Traversable f => f x -> f Int
indices f a
f) g Int
rhs
      | Bool
otherwise = Maybe (Entry f g)
forall a. Maybe a
Nothing

-- | Convert from 'NatMap'
fromNatMap :: (Traversable f, Traversable g) => NatMap f g -> PreNatMap f g
fromNatMap :: forall (f :: * -> *) (g :: * -> *).
(Traversable f, Traversable g) =>
NatMap f g -> PreNatMap f g
fromNatMap NatMap f g
nm = Map (Shape f) (PosData g) -> PreNatMap f g
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (PosData g) -> PreNatMap f g
PreNatMap (Map (Shape f) (PosData g) -> PreNatMap f g)
-> Map (Shape f) (PosData g) -> PreNatMap f g
forall a b. (a -> b) -> a -> b
$ [(Shape f, PosData g)] -> Map (Shape f) (PosData g)
forall k a. [(k, a)] -> Map k a
Map.fromDistinctAscList ([(Shape f, PosData g)] -> Map (Shape f) (PosData g))
-> [(Shape f, PosData g)] -> Map (Shape f) (PosData g)
forall a b. (a -> b) -> a -> b
$
  Entry f g -> (Shape f, PosData g)
forall {g :: * -> *} {f :: * -> *}.
(Traversable g, Foldable f) =>
Entry f g -> (Shape f, PosData g)
entryToPosData (Entry f g -> (Shape f, PosData g))
-> [Entry f g] -> [(Shape f, PosData g)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NatMap f g -> [Entry f g]
forall (f :: * -> *) (g :: * -> *). NatMap f g -> [Entry f g]
NM.toEntries NatMap f g
nm
    where
      entryToPosData :: Entry f g -> (Shape f, PosData g)
entryToPosData Entry f g
e = (Shape f
f, PosData g
pd)
        where
          (Shape f
f, g Var
gx) = Entry f g -> (Shape f, g Var)
forall (f :: * -> *) (g :: * -> *). Entry f g -> (Shape f, g Var)
NM.getKeyValue Entry f g
e
          vars :: [Var]
vars = Int -> [Var]
NM.makeVars (Shape f -> Int
forall (f :: * -> *). Foldable f => Shape f -> Int
lengthShape Shape f
f)
          unreachable :: a
unreachable = String -> a
forall a. HasCallStack => String -> a
error String
"this makePosData can't fail"
          pd :: PosData g
pd = PosData g -> Maybe (PosData g) -> PosData g
forall a. a -> Maybe a -> a
fromMaybe PosData g
forall {a}. a
unreachable (Maybe (PosData g) -> PosData g) -> Maybe (PosData g) -> PosData g
forall a b. (a -> b) -> a -> b
$ [Var] -> g Var -> Maybe (PosData g)
forall a (g :: * -> *).
(Ord a, Traversable g) =>
[a] -> g a -> Maybe (PosData g)
makePosData [Var]
vars g Var
gx

-- | Extract only the 'Shape' part as a plain 'Map'.
toShapeMap :: PreNatMap f g -> Map (Shape f) (Shape g)
toShapeMap :: forall (f :: * -> *) (g :: * -> *).
PreNatMap f g -> Map (Shape f) (Shape g)
toShapeMap (PreNatMap Map (Shape f) (PosData g)
pnm) = (PosData g -> Shape g)
-> Map (Shape f) (PosData g) -> Map (Shape f) (Shape g)
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map (\ (PosData Vector Int
_ g Int
gi) -> g Int -> Shape g
forall (f :: * -> *) a. f a -> Shape f
Shape g Int
gi) Map (Shape f) (PosData g)
pnm

-- | Convert the 'Map' of 'Shape' parts to @PreNatMap@, by each mapping
--   of shape @f0 :: Shape f@ to shape @g0 :: Shape g@ as an example input-output
--   pair of contentless values @(f (), g ())@.
--
--   Note that this operation can fail when @f@ is empty and @g@ is nonempty.
--   For example, a @Map@ containing mapping @Shape Nothing@ to @Shape (Just ())@
--   fails, as no natural transformation @Maybe ~> Maybe@ can turn @Nothing@ to @Just x@.
fromShapeMap :: (Foldable f, Traversable g) => Map (Shape f) (Shape g) -> Maybe (PreNatMap f g)
fromShapeMap :: forall (f :: * -> *) (g :: * -> *).
(Foldable f, Traversable g) =>
Map (Shape f) (Shape g) -> Maybe (PreNatMap f g)
fromShapeMap Map (Shape f) (Shape g)
m = Map (Shape f) (PosData g) -> PreNatMap f g
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (PosData g) -> PreNatMap f g
PreNatMap (Map (Shape f) (PosData g) -> PreNatMap f g)
-> Maybe (Map (Shape f) (PosData g)) -> Maybe (PreNatMap f g)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map (Shape f) (Maybe (PosData g))
-> Maybe (Map (Shape f) (PosData g))
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: * -> *) a.
Applicative f =>
Map (Shape f) (f a) -> f (Map (Shape f) a)
sequenceA ((Shape f -> Shape g -> Maybe (PosData g))
-> Map (Shape f) (Shape g) -> Map (Shape f) (Maybe (PosData g))
forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey Shape f -> Shape g -> Maybe (PosData g)
forall (f :: * -> *) (g :: * -> *).
(Foldable f, Traversable g) =>
Shape f -> Shape g -> Maybe (PosData g)
makeShapePosData Map (Shape f) (Shape g)
m)

-- | Query the output of natural transformation for given input.
--   Succeeds only when the output is uniquely determined.
match :: (Eq a, Ord (f Ignored), Foldable f, Functor g) => f a -> PreNatMap f g -> Maybe (g a)
match :: forall a (f :: * -> *) (g :: * -> *).
(Eq a, Ord (f Ignored), Foldable f, Functor g) =>
f a -> PreNatMap f g -> Maybe (g a)
match f a
fa (PreNatMap Map (Shape f) (PosData g)
pnm) = do
  pd@(PosData lhs rhs) <- Shape f -> Map (Shape f) (PosData g) -> Maybe (PosData g)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (f a -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f a
fa) Map (Shape f) (PosData g)
pnm
  if isCompleteLHS pd
    then pure $ makeRHS (toList fa) pd
    else do
      funLhs <- functionalRelI (zip (UV.toList lhs) (toList fa))
      let ga = (IntMap a
funLhs IntMap a -> Int -> a
forall a. IntMap a -> Int -> a
IntMap.!) (Int -> a) -> g Int -> g a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g Int
rhs
      pure ga
{-# INLINABLE match #-}

-- | Query the output of natural transformation for given input @fa :: f a@.
--   Succeeds only when the @PreNatMap@ had full knowledge for the inputs
--   with same shape to @fa@.
--
-- Compared to @match@, @fullMatch@ does not require @Eq a@ constraint
-- for the type of contents, by giving up cases which requires checks for
-- equality comparison between contents of the input @fa@. 
fullMatch :: (Ord (f Ignored), Foldable f, Functor g) => f a -> PreNatMap f g -> Maybe (g a)
fullMatch :: forall (f :: * -> *) (g :: * -> *) a.
(Ord (f Ignored), Foldable f, Functor g) =>
f a -> PreNatMap f g -> Maybe (g a)
fullMatch f a
fa (PreNatMap Map (Shape f) (PosData g)
pnm) = do
  pd <- Shape f -> Map (Shape f) (PosData g) -> Maybe (PosData g)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (f a -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f a
fa) Map (Shape f) (PosData g)
pnm
  if isCompleteLHS pd
    then pure $ makeRHS (toList fa) pd
    else Nothing
{-# INLINABLE fullMatch #-}

isFull :: (Ord (f Ignored), Foldable f, Functor g) => Shape f -> PreNatMap f g -> Bool
isFull :: forall (f :: * -> *) (g :: * -> *).
(Ord (f Ignored), Foldable f, Functor g) =>
Shape f -> PreNatMap f g -> Bool
isFull Shape f
f (PreNatMap Map (Shape f) (PosData g)
pnm) = Bool -> (PosData g -> Bool) -> Maybe (PosData g) -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False PosData g -> Bool
forall (g :: * -> *). PosData g -> Bool
isCompleteLHS (Maybe (PosData g) -> Bool) -> Maybe (PosData g) -> Bool
forall a b. (a -> b) -> a -> b
$ Shape f -> Map (Shape f) (PosData g) -> Maybe (PosData g)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Shape f
f Map (Shape f) (PosData g)
pnm
{-# INLINABLE isFull #-}

-- | Query the output of natural transformation for given input @fa :: f a@.
--   Succeeds if the shape of the output corresponding to the input
--   is known.
--
--   'lookup' succeeds even if the output is not determined uniquely.
--   If there are multiple contents of inputs for a content of the output,
--   'lookup' merges all of the candidate contents using 'Semigroup' operation.
lookup :: (Semigroup a, Ord (f Ignored), Foldable f, Functor g) => f a -> PreNatMap f g -> Maybe (g a)
lookup :: forall a (f :: * -> *) (g :: * -> *).
(Semigroup a, Ord (f Ignored), Foldable f, Functor g) =>
f a -> PreNatMap f g -> Maybe (g a)
lookup = (a -> a) -> f a -> PreNatMap f g -> Maybe (g a)
forall b (f :: * -> *) (g :: * -> *) a.
(Semigroup b, Ord (f Ignored), Foldable f, Functor g) =>
(a -> b) -> f a -> PreNatMap f g -> Maybe (g b)
lookupWith a -> a
forall a. a -> a
id
{-# INLINABLE lookup #-}

-- | Query the output of natural transformation for given input @fa :: f a@,
--   while mapping its contents to another type @b@ with a @Semigroup@ instance.
--
-- > lookupWith h fa === lookup (h <$> fa)
lookupWith :: (Semigroup b, Ord (f Ignored), Foldable f, Functor g) => (a -> b) -> f a -> PreNatMap f g -> Maybe (g b)
lookupWith :: forall b (f :: * -> *) (g :: * -> *) a.
(Semigroup b, Ord (f Ignored), Foldable f, Functor g) =>
(a -> b) -> f a -> PreNatMap f g -> Maybe (g b)
lookupWith a -> b
h f a
fa (PreNatMap Map (Shape f) (PosData g)
pnm) = do
  pd@(PosData lhs rhs) <- Shape f -> Map (Shape f) (PosData g) -> Maybe (PosData g)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (f a -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f a
fa) Map (Shape f) (PosData g)
pnm
  let bs = a -> b
h (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f a
fa
      funLhs = (b -> b -> b) -> [(Int, b)] -> IntMap b
forall a. (a -> a -> a) -> [(Int, a)] -> IntMap a
IntMap.fromListWith b -> b -> b
forall a. Semigroup a => a -> a -> a
(<>) ([Int] -> [b] -> [(Int, b)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
UV.toList Vector Int
lhs) [b]
bs)
      gb | PosData g -> Bool
forall (g :: * -> *). PosData g -> Bool
isCompleteLHS PosData g
pd = [b] -> PosData g -> g b
forall (g :: * -> *) a. Functor g => [a] -> PosData g -> g a
makeRHS [b]
bs PosData g
pd
         | Bool
otherwise = (IntMap b
funLhs IntMap b -> Int -> b
forall a. IntMap a -> Int -> a
IntMap.!) (Int -> b) -> g Int -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g Int
rhs
  pure gb
{-# INLINABLE lookupWith #-}

-- | Looks up only the shape.
--
-- @
-- lookupShape (Shape f) p === Shape '<$>' lookup ('Data.Functor.void' f) p
-- @
lookupShape :: Ord (f Ignored) => Shape f -> PreNatMap f g -> Maybe (Shape g)
lookupShape :: forall (f :: * -> *) (g :: * -> *).
Ord (f Ignored) =>
Shape f -> PreNatMap f g -> Maybe (Shape g)
lookupShape Shape f
f (PreNatMap Map (Shape f) (PosData g)
pnm) = case Shape f -> Map (Shape f) (PosData g) -> Maybe (PosData g)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Shape f
f Map (Shape f) (PosData g)
pnm of
  Maybe (PosData g)
Nothing -> Maybe (Shape g)
forall a. Maybe a
Nothing
  Just (PosData Vector Int
_ g Int
rhs) -> Shape g -> Maybe (Shape g)
forall a. a -> Maybe a
Just (g Int -> Shape g
forall (f :: * -> *) a. f a -> Shape f
Shape g Int
rhs)
{-# INLINABLE lookupShape #-}

-- | Refines knowledge a 'PreNatMap' contains by a pair of example
--   input-output pair.
--
-- Returns @Nothing@ if the given example is not consistent with the already given
-- examples.
refine :: (Ord a, Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g) => f a -> g a -> PreNatMap f g -> Maybe (PreNatMap f g)
refine :: forall a (f :: * -> *) (g :: * -> *).
(Ord a, Ord (f Ignored), Eq (g Ignored), Foldable f,
 Traversable g) =>
f a -> g a -> PreNatMap f g -> Maybe (PreNatMap f g)
refine f a
fa g a
ga = ((Bool, PreNatMap f g) -> PreNatMap f g)
-> Maybe (Bool, PreNatMap f g) -> Maybe (PreNatMap f g)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool, PreNatMap f g) -> PreNatMap f g
forall a b. (a, b) -> b
snd (Maybe (Bool, PreNatMap f g) -> Maybe (PreNatMap f g))
-> (PreNatMap f g -> Maybe (Bool, PreNatMap f g))
-> PreNatMap f g
-> Maybe (PreNatMap f g)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a -> g a -> PreNatMap f g -> Maybe (Bool, PreNatMap f g)
forall a (f :: * -> *) (g :: * -> *).
(Ord a, Ord (f Ignored), Eq (g Ignored), Foldable f,
 Traversable g) =>
f a -> g a -> PreNatMap f g -> Maybe (Bool, PreNatMap f g)
refine' f a
fa g a
ga
{-# INLINABLE refine #-}

-- | Same to 'refine' but also returns if any change has been made.
refine' :: (Ord a, Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g) => f a -> g a -> PreNatMap f g -> Maybe (Bool, PreNatMap f g)
refine' :: forall a (f :: * -> *) (g :: * -> *).
(Ord a, Ord (f Ignored), Eq (g Ignored), Foldable f,
 Traversable g) =>
f a -> g a -> PreNatMap f g -> Maybe (Bool, PreNatMap f g)
refine' f a
fa g a
ga (PreNatMap Map (Shape f) (PosData g)
pnm) = case Shape f -> Map (Shape f) (PosData g) -> Maybe (PosData g)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (f a -> Shape f
forall (f :: * -> *) a. f a -> Shape f
Shape f a
fa) Map (Shape f) (PosData g)
pnm of
  -- Create new entry
  Maybe (PosData g)
Nothing -> do
    pd <- [a] -> g a -> Maybe (PosData g)
forall a (g :: * -> *).
(Ord a, Traversable g) =>
[a] -> g a -> Maybe (PosData g)
makePosData (f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f a
fa) g a
ga
    pure (True, PreNatMap $ Map.insert (Shape fa) pd pnm)
  -- Update exising entry
  Just PosData g
pd -> do
    (changed, pd') <- PosData g -> [a] -> g a -> Maybe (Bool, PosData g)
forall a (g :: * -> *).
(Ord a, Eq (g Ignored), Traversable g) =>
PosData g -> [a] -> g a -> Maybe (Bool, PosData g)
refinePosData PosData g
pd (f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f a
fa) g a
ga
    pure (changed, PreNatMap $ Map.insert (Shape fa) pd' pnm)
{-# INLINABLE refine' #-}

-- | 'refine' but only the shapes of the input and output is given.
refineShape :: (Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g)
  => Shape f -> Shape g -> PreNatMap f g -> Maybe (PreNatMap f g)
refineShape :: forall (f :: * -> *) (g :: * -> *).
(Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g) =>
Shape f -> Shape g -> PreNatMap f g -> Maybe (PreNatMap f g)
refineShape Shape f
f Shape g
g = ((Bool, PreNatMap f g) -> PreNatMap f g)
-> Maybe (Bool, PreNatMap f g) -> Maybe (PreNatMap f g)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool, PreNatMap f g) -> PreNatMap f g
forall a b. (a, b) -> b
snd (Maybe (Bool, PreNatMap f g) -> Maybe (PreNatMap f g))
-> (PreNatMap f g -> Maybe (Bool, PreNatMap f g))
-> PreNatMap f g
-> Maybe (PreNatMap f g)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Shape f -> Shape g -> PreNatMap f g -> Maybe (Bool, PreNatMap f g)
forall (f :: * -> *) (g :: * -> *).
(Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g) =>
Shape f -> Shape g -> PreNatMap f g -> Maybe (Bool, PreNatMap f g)
refineShape' Shape f
f Shape g
g
{-# INLINABLE refineShape #-}

-- | Same to 'refineShape' but also returns if any change has been made.
refineShape' :: (Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g)
  => Shape f -> Shape g -> PreNatMap f g -> Maybe (Bool, PreNatMap f g)
refineShape' :: forall (f :: * -> *) (g :: * -> *).
(Ord (f Ignored), Eq (g Ignored), Foldable f, Traversable g) =>
Shape f -> Shape g -> PreNatMap f g -> Maybe (Bool, PreNatMap f g)
refineShape' Shape f
f Shape g
g (PreNatMap Map (Shape f) (PosData g)
pnm) = case Shape f -> Map (Shape f) (PosData g) -> Maybe (PosData g)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Shape f
f Map (Shape f) (PosData g)
pnm of
  -- Create new entry
  Maybe (PosData g)
Nothing -> do
    pd <- Shape f -> Shape g -> Maybe (PosData g)
forall (f :: * -> *) (g :: * -> *).
(Foldable f, Traversable g) =>
Shape f -> Shape g -> Maybe (PosData g)
makeShapePosData Shape f
f Shape g
g
    pure (True, PreNatMap $ Map.insert f pd pnm)
  -- Check existing entry has the matcing shape 
  Just (PosData Vector Int
_ g Int
gx) -> (Bool
False, Map (Shape f) (PosData g) -> PreNatMap f g
forall (f :: * -> *) (g :: * -> *).
Map (Shape f) (PosData g) -> PreNatMap f g
PreNatMap Map (Shape f) (PosData g)
pnm) (Bool, PreNatMap f g) -> Maybe () -> Maybe (Bool, PreNatMap f g)
forall a b. a -> Maybe b -> Maybe a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Shape g
g Shape g -> Shape g -> Bool
forall a. Eq a => a -> a -> Bool
== g Int -> Shape g
forall (f :: * -> *) a. f a -> Shape f
Shape g Int
gx)
{-# INLINABLE refineShape' #-}

---- Utilities ----

-- | Tests if the "lhs" vector of @PosData@ contains
--   all distinct values.
--
--   By the invariant of @PosData@, it is not necessary
--   to check every element.
isCompleteLHS :: PosData g -> Bool
isCompleteLHS :: forall (g :: * -> *). PosData g -> Bool
isCompleteLHS (PosData Vector Int
lhs g Int
_) = Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Vector Int
lhs Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
UV.! (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
UV.length Vector Int
lhs

-- | Creates RHS ignoring LHS variable pattern,
--   assuming 'isCompleteLHS' hold.
makeRHS :: Functor g => [a] -> PosData g -> g a
makeRHS :: forall (g :: * -> *) a. Functor g => [a] -> PosData g -> g a
makeRHS [a]
as (PosData Vector Int
_ g Int
gi) =
  let aVec :: Vector a
aVec = [a] -> Vector a
forall a. [a] -> Vector a
V.fromList [a]
as
  in (Vector a
aVec Vector a -> Int -> a
forall a. Vector a -> Int -> a
V.!) (Int -> a) -> g Int -> g a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> g Int
gi

-- Create an IntMap if the relation is functional
functionalRelI :: (Eq a) => [(Int, a)] -> Maybe (IntMap.IntMap a)
functionalRelI :: forall a. Eq a => [(Int, a)] -> Maybe (IntMap a)
functionalRelI = IntMap a -> [(Int, a)] -> Maybe (IntMap a)
forall {a}. Eq a => IntMap a -> [(Int, a)] -> Maybe (IntMap a)
loop IntMap a
forall a. IntMap a
IntMap.empty
  where
    loop :: IntMap a -> [(Int, a)] -> Maybe (IntMap a)
loop !IntMap a
m [] = IntMap a -> Maybe (IntMap a)
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntMap a
m
    loop !IntMap a
m ((Int
k,a
v) : [(Int, a)]
rest) = case Int -> IntMap a -> Maybe a
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
k IntMap a
m of
      Maybe a
Nothing -> IntMap a -> [(Int, a)] -> Maybe (IntMap a)
loop (Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
k a
v IntMap a
m) [(Int, a)]
rest
      Just a
v'
        | a
v a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
v' -> IntMap a -> [(Int, a)] -> Maybe (IntMap a)
loop IntMap a
m [(Int, a)]
rest
        | Bool
otherwise -> Maybe (IntMap a)
forall a. Maybe a
Nothing

makeShapePosData :: (Foldable f, Traversable g)
  => Shape f -> Shape g -> Maybe (PosData g)
makeShapePosData :: forall (f :: * -> *) (g :: * -> *).
(Foldable f, Traversable g) =>
Shape f -> Shape g -> Maybe (PosData g)
makeShapePosData (Shape f a
f) (Shape g a
g)
  | f a -> Bool
forall a. f a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null f a
f Bool -> Bool -> Bool
&& Bool -> Bool
not (g a -> Bool
forall a. g a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null g a
g) = Maybe (PosData g)
forall a. Maybe a
Nothing
  | Bool
otherwise = PosData g -> Maybe (PosData g)
forall a. a -> Maybe a
Just (PosData g -> Maybe (PosData g)) -> PosData g -> Maybe (PosData g)
forall a b. (a -> b) -> a -> b
$ Vector Int -> g Int -> PosData g
forall (g :: * -> *). Vector Int -> g Int -> PosData g
PosData Vector Int
lhs g Int
rhs
  where
    lhs :: Vector Int
lhs = Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
UV.replicate (f a -> Int
forall a. f a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length f a
f) Int
0
    rhs :: g Int
rhs = Int
0 Int -> g a -> g Int
forall a b. a -> g b -> g a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ g a
g

-- >>> makePosData "foo" "bar"
-- Nothing
-- >>> makePosData "banana" "anba"
-- Just (PosData [0,1,2,1,2,1] [1,2,0,1])
-- >>> makePosData "banana" (Just 'b')
-- Just (PosData [0,1,2,3,4,5] (Just 0))
makePosData :: (Ord a, Traversable g) => [a] -> g a -> Maybe (PosData g)
makePosData :: forall a (g :: * -> *).
(Ord a, Traversable g) =>
[a] -> g a -> Maybe (PosData g)
makePosData [a]
as g a
ga = Vector Int -> g Int -> PosData g
forall (g :: * -> *). Vector Int -> g Int -> PosData g
PosData Vector Int
lhs (g Int -> PosData g) -> Maybe (g Int) -> Maybe (PosData g)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (g Int)
rhs
  where
    rhsSet :: Set a
rhsSet = [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList (g a -> [a]
forall a. g a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList g a
ga)
    ([Int]
lhsList, Map a Int
revmap) = (a -> Bool) -> [a] -> ([Int], Map a Int)
forall a. Ord a => (a -> Bool) -> [a] -> ([Int], Map a Int)
reindex (a -> Set a -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set a
rhsSet) [a]
as
    lhs :: Vector Int
lhs = [Int] -> Vector Int
forall a. Unbox a => [a] -> Vector a
UV.fromList [Int]
lhsList
    rhs :: Maybe (g Int)
rhs = (a -> Maybe Int) -> g a -> Maybe (g Int)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> g a -> f (g b)
traverse (a -> Map a Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
`Map.lookup` Map a Int
revmap) g a
ga

refinePosData :: (Ord a, Eq (g Ignored), Traversable g) => PosData g -> [a] -> g a -> Maybe (Bool, PosData g)
refinePosData :: forall a (g :: * -> *).
(Ord a, Eq (g Ignored), Traversable g) =>
PosData g -> [a] -> g a -> Maybe (Bool, PosData g)
refinePosData (PosData Vector Int
lhs g Int
rhs) [a]
as g a
ga = do
  ga' <- g Int -> g a -> Maybe (g (Int, a))
forall (t :: * -> *) a b.
(Eq (t Ignored), Traversable t) =>
t a -> t b -> Maybe (t (a, b))
zipMatch g Int
rhs g a
ga
  pd@(PosData lhs' _) <- makePosData (zip (UV.toList lhs) as) ga'
  guard $ UV.length lhs == UV.length lhs'
  pure (lhs /= lhs', pd)

reindex :: (Ord a) => (a -> Bool) -> [a] -> ([Int], Map a Int)
reindex :: forall a. Ord a => (a -> Bool) -> [a] -> ([Int], Map a Int)
reindex a -> Bool
mask = Int -> Map a Int -> [a] -> ([Int], Map a Int)
forall {a}. Num a => a -> Map a a -> [a] -> ([a], Map a a)
loop Int
0 Map a Int
forall k a. Map k a
Map.empty
  where
    loop :: a -> Map a a -> [a] -> ([a], Map a a)
loop a
_ Map a a
rev [] = ([], Map a a
rev)
    loop !a
n !Map a a
rev (a
a : [a]
rest)
      | a -> Bool
mask a
a = case a -> Map a a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup a
a Map a a
rev of
          Maybe a
Nothing -> case a -> Map a a -> [a] -> ([a], Map a a)
loop (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) (a -> a -> Map a a -> Map a a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert a
a a
n Map a a
rev) [a]
rest of
            ~([a]
ks,Map a a
rev') -> (a
n a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ks, Map a a
rev')
          Just a
k  -> case a -> Map a a -> [a] -> ([a], Map a a)
loop a
n Map a a
rev [a]
rest of
            ~([a]
ks,Map a a
rev') -> (a
k a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ks, Map a a
rev')
      | Bool
otherwise = case a -> Map a a -> [a] -> ([a], Map a a)
loop (a
n a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) Map a a
rev [a]
rest of
            ~([a]
ks,Map a a
rev') -> (a
n a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ks, Map a a
rev')