{-# Language ImportQualifiedPost, RoleAnnotations, DataKinds, KindSignatures #-}
{-|
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
import Data.Vector.Unboxed qualified as V
import GHC.TypeLits

-- $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, ReadPrec [Permutation n]
ReadPrec (Permutation n)
Int -> ReadS (Permutation n)
ReadS [Permutation n]
(Int -> ReadS (Permutation n))
-> ReadS [Permutation n]
-> ReadPrec (Permutation n)
-> ReadPrec [Permutation n]
-> Read (Permutation n)
forall (n :: Nat). ReadPrec [Permutation n]
forall (n :: Nat). ReadPrec (Permutation n)
forall (n :: Nat). Int -> ReadS (Permutation n)
forall (n :: Nat). ReadS [Permutation n]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: forall (n :: Nat). Int -> ReadS (Permutation n)
readsPrec :: Int -> ReadS (Permutation n)
$creadList :: forall (n :: Nat). ReadS [Permutation n]
readList :: ReadS [Permutation n]
$creadPrec :: forall (n :: Nat). ReadPrec (Permutation n)
readPrec :: ReadPrec (Permutation n)
$creadListPrec :: forall (n :: Nat). ReadPrec [Permutation n]
readListPrec :: ReadPrec [Permutation n]
Read, Int -> Permutation n -> ShowS
[Permutation n] -> ShowS
Permutation n -> String
(Int -> Permutation n -> ShowS)
-> (Permutation n -> String)
-> ([Permutation n] -> ShowS)
-> Show (Permutation n)
forall (n :: Nat). Int -> Permutation n -> ShowS
forall (n :: Nat). [Permutation n] -> ShowS
forall (n :: Nat). Permutation n -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall (n :: Nat). Int -> Permutation n -> ShowS
showsPrec :: Int -> Permutation n -> ShowS
$cshow :: forall (n :: Nat). Permutation n -> String
show :: Permutation n -> String
$cshowList :: forall (n :: Nat). [Permutation n] -> ShowS
showList :: [Permutation n] -> ShowS
Show)

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

-- | Size of the list of elements permuted.
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 :: Permutation n -> Bool
isValid :: forall (n :: Nat). Permutation n -> Bool
isValid (P Vector Int
p) = Vector Bool -> Bool
V.and ((Bool -> Bool -> Bool)
-> Vector Bool -> Vector Int -> Vector Bool -> Vector Bool
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> Vector a -> Vector Int -> Vector b -> Vector a
V.accumulate_ (\Bool
_ Bool
new -> Bool
new) (Int -> Bool -> Vector Bool
forall a. Unbox a => Int -> a -> Vector a
V.replicate Int
n Bool
False) Vector Int
p (Int -> Bool -> Vector Bool
forall a. Unbox a => Int -> a -> Vector a
V.replicate Int
n Bool
True))
  where
    n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
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 -> Permutation n) -> Permutation n)
-> (Int -> Permutation n) -> Permutation n
forall a b. (a -> b) -> a -> b
$ \Int
n -> Vector Int -> Permutation n
forall (n :: Nat). Vector Int -> Permutation n
P (Vector Int -> Permutation n) -> Vector Int -> Permutation n
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
V.generate Int
n ((Int -> Int) -> Vector Int) -> (Int -> Int) -> Vector Int
forall a b. (a -> b) -> a -> b
$ \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 :: 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 -> Permutation n) -> Permutation n)
-> (Int -> Permutation n) -> Permutation n
forall a b. (a -> b) -> a -> b
$ \Int
n ->
  let x' :: Int
x' = Int
xInt -> Int -> Int
forall a. Integral a => a -> a -> a
`mod`Int
n -- not evaluated when n == 0
      y' :: Int
y' = Int
yInt -> 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 -> Int) -> Permutation n) -> (Int -> Int) -> Permutation n
forall a b. (a -> b) -> a -> b
$ \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 :: KnownNat n => Int -> Permutation n
rotateRight :: forall (n :: Nat). KnownNat n => Int -> Permutation n
rotateRight = Int -> Permutation n
forall (n :: Nat). KnownNat n => Int -> Permutation n
rotateLeft (Int -> Permutation n) -> (Int -> Int) -> Int -> Permutation n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int
forall a. Num a => a -> a
negate

-- | Permutation generated by rotating all the elements to the left.
rotateLeft :: KnownNat n => Int -> Permutation n
rotateLeft :: forall (n :: Nat). KnownNat n => Int -> Permutation n
rotateLeft Int
n = (Int -> Int) -> Permutation n
forall (n :: Nat). KnownNat n => (Int -> Int) -> Permutation n
mkPermutation ((Int -> Int) -> Permutation n) -> (Int -> Int) -> Permutation n
forall a b. (a -> b) -> a -> b
$ \Int
i -> Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
n

-- | Permutation generated by inverting another permutation.
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 ((Int -> Int -> Int)
-> Vector Int -> Vector Int -> Vector Int -> Vector Int
forall a b.
(Unbox a, Unbox b) =>
(a -> b -> a) -> Vector a -> Vector Int -> Vector b -> Vector a
V.accumulate_ (\Int
_ Int
new -> Int
new) 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 :: 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 -> Int) -> Permutation n) -> (Int -> Int) -> Permutation n
forall a b. (a -> b) -> a -> b
$ \Int
i -> -Int
iInt -> 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 (P Vector Int
v) = (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 .. Vector Int -> Int
forall a. Unbox a => Vector a -> Int
V.length Vector Int
vInt -> 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 (Vector Int
v Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
V.! Int
cur)

    aux :: IntSet -> Maybe ([Int], IntSet)
aux IntSet
items =
      do (Int
seed, IntSet
items') <- IntSet -> Maybe (Int, IntSet)
IntSet.minView IntSet
items
         let cycleElts :: [Int]
cycleElts = Int
seed Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int -> Int -> [Int]
getOne Int
seed (Vector Int
v Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
V.! Int
seed)
         ([Int], IntSet) -> Maybe ([Int], IntSet)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Int]
cycleElts, (IntSet -> Int -> IntSet) -> IntSet -> [Int] -> IntSet
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((Int -> IntSet -> IntSet) -> IntSet -> Int -> IntSet
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> IntSet -> IntSet
IntSet.delete) IntSet
items' [Int]
cycleElts)

-- | 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@.
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

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)

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