{-# Language ImportQualifiedPost, RoleAnnotations, DataKinds, KindSignatures, BlockArguments, ScopedTypeVariables #-}
{-|
Module      : Advent.Permutation
Description : Composable permutations
Copyright   : (c) Eric Mertens, 2017
License     : ISC
Maintainer  : emertens@gmail.com

Common permutations of a finite collection of elements and
operations on them.

-}
module Advent.Permutation
  ( Permutation
  , runPermutation
  , mkPermutation
  , swap
  , rotateRight
  , rotateLeft
  , invert
  , isValid
  , size
  , backwards
  , cycles
  , order
  ) where

import Advent.Group (Group(..))
import Data.Function (fix)
import Data.IntSet qualified as IntSet
import Data.List (foldl', unfoldr)
import Data.List.NonEmpty (NonEmpty((:|)))
import Data.Semigroup (Semigroup(sconcat))
import Data.Vector.Unboxed qualified as V
import GHC.TypeLits (KnownNat, Nat, natVal)

-- $setup
-- >>> :set -XDataKinds

type role Permutation nominal

-- | Permutations of @n@ elements
newtype Permutation (n :: Nat) = P (V.Vector Int)
  deriving (Permutation n -> Permutation n -> Bool
(Permutation n -> Permutation n -> Bool)
-> (Permutation n -> Permutation n -> Bool) -> Eq (Permutation n)
forall (n :: Nat). Permutation n -> Permutation n -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (n :: Nat). Permutation n -> Permutation n -> Bool
== :: Permutation n -> Permutation n -> Bool
$c/= :: forall (n :: Nat). Permutation n -> Permutation n -> Bool
/= :: Permutation n -> Permutation n -> Bool
Eq, Eq (Permutation n)
Eq (Permutation n) =>
(Permutation n -> Permutation n -> Ordering)
-> (Permutation n -> Permutation n -> Bool)
-> (Permutation n -> Permutation n -> Bool)
-> (Permutation n -> Permutation n -> Bool)
-> (Permutation n -> Permutation n -> Bool)
-> (Permutation n -> Permutation n -> Permutation n)
-> (Permutation n -> Permutation n -> Permutation n)
-> Ord (Permutation n)
Permutation n -> Permutation n -> Bool
Permutation n -> Permutation n -> Ordering
Permutation n -> Permutation n -> Permutation n
forall (n :: Nat). Eq (Permutation n)
forall (n :: Nat). Permutation n -> Permutation n -> Bool
forall (n :: Nat). Permutation n -> Permutation n -> Ordering
forall (n :: Nat). Permutation n -> Permutation n -> Permutation n
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: forall (n :: Nat). Permutation n -> Permutation n -> Ordering
compare :: Permutation n -> Permutation n -> Ordering
$c< :: forall (n :: Nat). Permutation n -> Permutation n -> Bool
< :: Permutation n -> Permutation n -> Bool
$c<= :: forall (n :: Nat). Permutation n -> Permutation n -> Bool
<= :: Permutation n -> Permutation n -> Bool
$c> :: forall (n :: Nat). Permutation n -> Permutation n -> Bool
> :: Permutation n -> Permutation n -> Bool
$c>= :: forall (n :: Nat). Permutation n -> Permutation n -> Bool
>= :: Permutation n -> Permutation n -> Bool
$cmax :: forall (n :: Nat). Permutation n -> Permutation n -> Permutation n
max :: Permutation n -> Permutation n -> Permutation n
$cmin :: forall (n :: Nat). Permutation n -> Permutation n -> Permutation n
min :: Permutation n -> Permutation n -> Permutation n
Ord)

-- | Produce a list mapping list elements from their indexes to their output elements.
runPermutation :: (Int -> a) -> Permutation n -> [a]
runPermutation :: forall a (n :: Nat). (Int -> a) -> Permutation n -> [a]
runPermutation Int -> a
f (P Vector Int
v) = Int -> a
f (Int -> a) -> [Int] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Vector Int -> [Int]
forall a. Unbox a => Vector a -> [a]
V.toList Vector Int
v
{-# Inline runPermutation #-}

-- | Size of the list of elements permuted.
--
-- >>> size (mempty :: Permutation 5)
-- 5
size :: Permutation n -> Int
size :: forall (n :: Nat). Permutation n -> Int
size (P Vector Int
v) = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
v

-- | Validate a permutation. A valid permutation will map each element in the input
-- to a unique element in the output.
--
-- >>> isValid (mempty :: Permutation 5)
-- True
--
-- >>> isValid (mkPermutation (const 0) :: Permutation 5)
-- False
isValid :: KnownNat n => Permutation n -> Bool
isValid :: forall (n :: Nat). KnownNat n => Permutation n -> Bool
isValid p :: Permutation n
p@(P Vector Int
v) =
  Permutation n -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal Permutation n
p Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Bool -> Bool -> Bool
&&
  Vector Bool -> Bool
V.and (Vector Bool -> Vector Int -> Vector Bool -> Vector Bool
forall a. Unbox a => Vector a -> Vector Int -> Vector a -> Vector a
V.update_ (Int -> Bool -> Vector Bool
forall a. Unbox a => Int -> a -> Vector a
V.replicate Int
n Bool
False) Vector Int
v (Int -> Bool -> Vector Bool
forall a. Unbox a => Int -> a -> Vector a
V.replicate Int
n Bool
True))
  where
    n :: Int
n = Permutation n -> Int
forall (n :: Nat). Permutation n -> Int
size Permutation n
p

-- | Helper function for making the size of a requested permutation available
-- while building the permutation.
withSize :: KnownNat n => (Int -> Permutation n) -> Permutation n
withSize :: forall (n :: Nat).
KnownNat n =>
(Int -> Permutation n) -> Permutation n
withSize Int -> Permutation n
f = (Permutation n -> Permutation n) -> Permutation n
forall a. (a -> a) -> a
fix (Int -> Permutation n
f (Int -> Permutation n)
-> (Permutation n -> Int) -> Permutation n -> Permutation n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Int)
-> (Permutation n -> Integer) -> Permutation n -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Permutation n -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal)

-- | Given a function mapping incoming indices to outgoing ones, construct
-- a new permutation value.
mkPermutation :: KnownNat n => (Int -> Int) -> Permutation n
mkPermutation :: forall (n :: Nat). KnownNat n => (Int -> Int) -> Permutation n
mkPermutation Int -> Int
f = (Int -> Permutation n) -> Permutation n
forall (n :: Nat).
KnownNat n =>
(Int -> Permutation n) -> Permutation n
withSize \Int
n -> Vector Int -> Permutation n
forall (n :: Nat). Vector Int -> Permutation n
P (Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
V.generate Int
n \Int
i -> Int -> Int
f Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n)

-- | Permutation generated by swapping the elements at a pair of indices.
--
-- >>> swap 2 3 :: Permutation 5
-- [0,1,3,2,4]
swap :: KnownNat n => Int -> Int -> Permutation n
swap :: forall (n :: Nat). KnownNat n => Int -> Int -> Permutation n
swap Int
x Int
y = (Int -> Permutation n) -> Permutation n
forall (n :: Nat).
KnownNat n =>
(Int -> Permutation n) -> Permutation n
withSize \Int
n ->
  let x' :: Int
x' = Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n -- not evaluated when n == 0
      y' :: Int
y' = Int
y Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
n
  in (Int -> Int) -> Permutation n
forall (n :: Nat). KnownNat n => (Int -> Int) -> Permutation n
mkPermutation \Int
i -> if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
x' then Int
y' else if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y' then Int
x' else Int
i

-- | Permutation generated by rotating all the elements to the right.
--
-- >>> rotateRight 2 :: Permutation 5
-- [3,4,0,1,2]
rotateRight :: KnownNat n => Int -> Permutation n
rotateRight :: forall (n :: Nat). KnownNat n => Int -> Permutation n
rotateRight = (Int -> Int) -> Permutation n
forall (n :: Nat). KnownNat n => (Int -> Int) -> Permutation n
mkPermutation ((Int -> Int) -> Permutation n)
-> (Int -> Int -> Int) -> Int -> Permutation n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract

-- | Permutation generated by rotating all the elements to the left.
--
-- >>> rotateLeft 2 :: Permutation 5
-- [2,3,4,0,1]
rotateLeft :: KnownNat n => Int -> Permutation n
rotateLeft :: forall (n :: Nat). KnownNat n => Int -> Permutation n
rotateLeft = (Int -> Int) -> Permutation n
forall (n :: Nat). KnownNat n => (Int -> Int) -> Permutation n
mkPermutation ((Int -> Int) -> Permutation n)
-> (Int -> Int -> Int) -> Int -> Permutation n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)

-- | Permutation generated by inverting another permutation.
--
-- >>> swap 1 2 <> swap 3 4 <> swap 4 5 :: Permutation 6
-- [0,2,1,4,5,3]
--
-- >>> invert (swap 1 2 <> swap 3 4 <> swap 4 5 :: Permutation 6)
-- [0,2,1,5,3,4]
invert :: Permutation n -> Permutation n
invert :: forall (n :: Nat). Permutation n -> Permutation n
invert (P Vector Int
v) = Vector Int -> Permutation n
forall (n :: Nat). Vector Int -> Permutation n
P (Vector Int -> Vector Int -> Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector Int -> Vector a -> Vector a
V.update_ Vector Int
initial Vector Int
v Vector Int
iota)
  where
    n :: Int
n       = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
v
    initial :: Vector Int
initial = Int -> Int -> Vector Int
forall a. Unbox a => Int -> a -> Vector a
V.replicate Int
n Int
0 -- 0 is arbitrary, should all be overwritten
    iota :: Vector Int
iota    = Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
V.generate Int
n Int -> Int
forall a. a -> a
id

-- | Permutation generated by reversing the order of the elements.
--
-- >>> backwards :: Permutation 4
-- [3,2,1,0]
backwards :: KnownNat n => Permutation n
backwards :: forall (n :: Nat). KnownNat n => Permutation n
backwards = (Int -> Int) -> Permutation n
forall (n :: Nat). KnownNat n => (Int -> Int) -> Permutation n
mkPermutation \Int
i -> -Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

-- | Compute the disjoint cycles of the permutation.
--
-- >>> cycles (swap 1 2 <> swap 3 4 <> swap 4 5 :: Permutation 6)
-- [[0],[1,2],[3,4,5]]
cycles :: Permutation n -> [[Int]]
cycles :: forall (n :: Nat). Permutation n -> [[Int]]
cycles Permutation n
p = (IntSet -> Maybe ([Int], IntSet)) -> IntSet -> [[Int]]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr IntSet -> Maybe ([Int], IntSet)
aux IntSet
initialSet
  where
    initialSet :: IntSet
initialSet = [Int] -> IntSet
IntSet.fromList [Int
0 .. Permutation n -> Int
forall (n :: Nat). Permutation n -> Int
size Permutation n
p Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]

    getOne :: Int -> Int -> [Int]
getOne Int
start Int
cur
      | Int
start Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
cur = []
      | Bool
otherwise    = Int
cur Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int -> Int -> [Int]
getOne Int
start (Permutation n -> Int -> Int
forall (n :: Nat). Permutation n -> Int -> Int
index Permutation n
p Int
cur)

    aux :: IntSet -> Maybe ([Int], IntSet)
aux IntSet
items =
     do (seed, items') <- IntSet -> Maybe (Int, IntSet)
IntSet.minView IntSet
items
        let cycleElts = Int
seed Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int -> Int -> [Int]
getOne Int
seed (Permutation n -> Int -> Int
forall (n :: Nat). Permutation n -> Int -> Int
index Permutation n
p Int
seed)
        pure (cycleElts, foldl' (flip IntSet.delete) items' cycleElts)

-- | Find the output element for the corresponding input element.
index :: Permutation n -> Int -> Int
index :: forall (n :: Nat). Permutation n -> Int -> Int
index (P Vector Int
p) Int
i = Vector Int
p Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
V.! Int
i

-- | Compute the order of a permutation.
--
-- >>> order (swap 1 2 <> swap 3 4 <> swap 4 5 :: Permutation 6)
-- 6
order :: Permutation n -> Int
order :: forall (n :: Nat). Permutation n -> Int
order = (Int -> Int -> Int) -> Int -> [Int] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Int -> Int -> Int
forall a. Integral a => a -> a -> a
lcm Int
1 ([Int] -> Int) -> (Permutation n -> [Int]) -> Permutation n -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([[Int]] -> [Int])
-> (Permutation n -> [[Int]]) -> Permutation n -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Permutation n -> [[Int]]
forall (n :: Nat). Permutation n -> [[Int]]
cycles

-- | @a '<>' b@ is the permutation that first permutes with @a@ and
-- then with @b@.
--
-- >>> swap 1 2 <> swap 3 4 <> swap 4 5 :: Permutation 6
-- [0,2,1,4,5,3]
instance Semigroup (Permutation n) where
  P Vector Int
x <> :: Permutation n -> Permutation n -> Permutation n
<> P Vector Int
y = Vector Int -> Permutation n
forall (n :: Nat). Vector Int -> Permutation n
P (Vector Int -> Vector Int -> Vector Int
forall a. Unbox a => Vector a -> Vector Int -> Vector a
V.backpermute Vector Int
x Vector Int
y)
  sconcat :: NonEmpty (Permutation n) -> Permutation n
sconcat (Permutation n
x :| [Permutation n]
xs) = (Permutation n -> Permutation n -> Permutation n)
-> Permutation n -> [Permutation n] -> Permutation n
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Permutation n -> Permutation n -> Permutation n
forall a. Semigroup a => a -> a -> a
(<>) Permutation n
x [Permutation n]
xs

-- | Extend the 'Semigroup' instance with an identity permutation as 'mempty'.
-- 
-- >>> mempty :: Permutation 6
-- [0,1,2,3,4,5]
instance KnownNat n => Monoid (Permutation n) where
  mempty :: Permutation n
mempty           = (Int -> Int) -> Permutation n
forall (n :: Nat). KnownNat n => (Int -> Int) -> Permutation n
mkPermutation Int -> Int
forall a. a -> a
id
  mconcat :: [Permutation n] -> Permutation n
mconcat []       = Permutation n
forall a. Monoid a => a
mempty
  mconcat (Permutation n
x : [Permutation n]
xs) = NonEmpty (Permutation n) -> Permutation n
forall a. Semigroup a => NonEmpty a -> a
sconcat (Permutation n
x Permutation n -> [Permutation n] -> NonEmpty (Permutation n)
forall a. a -> [a] -> NonEmpty a
:| [Permutation n]
xs)

-- | Extends the 'Monoid' instance using 'invert'
instance KnownNat n => Group (Permutation n) where
  inverse :: Permutation n -> Permutation n
inverse = Permutation n -> Permutation n
forall (n :: Nat). Permutation n -> Permutation n
invert

-- | Render a permutation as a list literal.
--
-- >>> show (mempty :: Permutation 4)
-- "[0,1,2,3]"
instance Show (Permutation n) where
  show :: Permutation n -> String
show (P Vector Int
p) = Vector Int -> String
forall a. Show a => a -> String
show Vector Int
p

-- | Parse a permutation as a list literal.
instance KnownNat n => Read (Permutation n) where
  readsPrec :: Int -> ReadS (Permutation n)
readsPrec Int
p String
str =
    [ (Permutation n
p, String
str')
    | (Vector Int
xs, String
str') <- Int -> ReadS (Vector Int)
forall a. Read a => Int -> ReadS a
readsPrec Int
p String
str
    , let p :: Permutation n
p = Vector Int -> Permutation n
forall (n :: Nat). Vector Int -> Permutation n
P Vector Int
xs :: Permutation n
    , Permutation n -> Bool
forall (n :: Nat). KnownNat n => Permutation n -> Bool
isValid Permutation n
p
    ]