{-# LANGUAGE TypeFamilies #-} module Data.Lattice where import Control.Monad import Control.Monad.Primitive import Data.Bits import Data.Coerce import qualified Data.List as L import qualified Data.Vector.Unboxed as U import qualified Data.Vector.Unboxed.Mutable as UM import Math.Prime import My.Prelude (rep, rev) infix 4 .<. class Poset a where (.<.) :: a -> a -> Bool zeta :: (Integral i) => a -> a -> i zeta a x a y | a x a -> a -> Bool forall a. Poset a => a -> a -> Bool .<. a y = i 1 | Bool otherwise = i 0 moebius :: (Integral i) => a -> a -> i class Lattice a where (/\) :: a -> a -> a (\/) :: a -> a -> a class FastZetaMoebius f where type Dim f fastZeta :: (Num a, U.Unbox a, PrimMonad m) => (Int -> f Int) -> Dim f -> UM.MVector (PrimState m) a -> m () fastMoebius :: (Num a, U.Unbox a, PrimMonad m) => (Int -> f Int) -> Dim f -> UM.MVector (PrimState m) a -> m () newtype NatOrd a = NatOrd {forall a. NatOrd a -> a getNatOrd :: a} deriving (NatOrd a -> NatOrd a -> Bool (NatOrd a -> NatOrd a -> Bool) -> (NatOrd a -> NatOrd a -> Bool) -> Eq (NatOrd a) forall a. Eq a => NatOrd a -> NatOrd a -> Bool forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a $c== :: forall a. Eq a => NatOrd a -> NatOrd a -> Bool == :: NatOrd a -> NatOrd a -> Bool $c/= :: forall a. Eq a => NatOrd a -> NatOrd a -> Bool /= :: NatOrd a -> NatOrd a -> Bool Eq, Eq (NatOrd a) Eq (NatOrd a) => (NatOrd a -> NatOrd a -> Ordering) -> (NatOrd a -> NatOrd a -> Bool) -> (NatOrd a -> NatOrd a -> Bool) -> (NatOrd a -> NatOrd a -> Bool) -> (NatOrd a -> NatOrd a -> Bool) -> (NatOrd a -> NatOrd a -> NatOrd a) -> (NatOrd a -> NatOrd a -> NatOrd a) -> Ord (NatOrd a) NatOrd a -> NatOrd a -> Bool NatOrd a -> NatOrd a -> Ordering NatOrd a -> NatOrd a -> NatOrd a 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 forall a. Ord a => Eq (NatOrd a) forall a. Ord a => NatOrd a -> NatOrd a -> Bool forall a. Ord a => NatOrd a -> NatOrd a -> Ordering forall a. Ord a => NatOrd a -> NatOrd a -> NatOrd a $ccompare :: forall a. Ord a => NatOrd a -> NatOrd a -> Ordering compare :: NatOrd a -> NatOrd a -> Ordering $c< :: forall a. Ord a => NatOrd a -> NatOrd a -> Bool < :: NatOrd a -> NatOrd a -> Bool $c<= :: forall a. Ord a => NatOrd a -> NatOrd a -> Bool <= :: NatOrd a -> NatOrd a -> Bool $c> :: forall a. Ord a => NatOrd a -> NatOrd a -> Bool > :: NatOrd a -> NatOrd a -> Bool $c>= :: forall a. Ord a => NatOrd a -> NatOrd a -> Bool >= :: NatOrd a -> NatOrd a -> Bool $cmax :: forall a. Ord a => NatOrd a -> NatOrd a -> NatOrd a max :: NatOrd a -> NatOrd a -> NatOrd a $cmin :: forall a. Ord a => NatOrd a -> NatOrd a -> NatOrd a min :: NatOrd a -> NatOrd a -> NatOrd a Ord, Int -> NatOrd a -> ShowS [NatOrd a] -> ShowS NatOrd a -> String (Int -> NatOrd a -> ShowS) -> (NatOrd a -> String) -> ([NatOrd a] -> ShowS) -> Show (NatOrd a) forall a. Show a => Int -> NatOrd a -> ShowS forall a. Show a => [NatOrd a] -> ShowS forall a. Show a => NatOrd a -> String forall a. (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a $cshowsPrec :: forall a. Show a => Int -> NatOrd a -> ShowS showsPrec :: Int -> NatOrd a -> ShowS $cshow :: forall a. Show a => NatOrd a -> String show :: NatOrd a -> String $cshowList :: forall a. Show a => [NatOrd a] -> ShowS showList :: [NatOrd a] -> ShowS Show) instance (Integral a, Ord a) => Poset (NatOrd a) where .<. :: NatOrd a -> NatOrd a -> Bool (.<.) = NatOrd a -> NatOrd a -> Bool forall a. Ord a => a -> a -> Bool (<=) {-# INLINE (.<.) #-} moebius :: forall i. Integral i => NatOrd a -> NatOrd a -> i moebius (NatOrd a x) (NatOrd a y) | a x a -> a -> Bool forall a. Eq a => a -> a -> Bool == a y = i 1 | a x a -> a -> a forall a. Num a => a -> a -> a + a 1 a -> a -> Bool forall a. Eq a => a -> a -> Bool == a y = -i 1 | Bool otherwise = i 0 instance (Ord a) => Lattice (NatOrd a) where /\ :: NatOrd a -> NatOrd a -> NatOrd a (/\) = (a -> a -> a) -> NatOrd a -> NatOrd a -> NatOrd a forall a b. Coercible a b => a -> b coerce (forall a. Ord a => a -> a -> a min @a) {-# INLINE (/\) #-} \/ :: NatOrd a -> NatOrd a -> NatOrd a (\/) = (a -> a -> a) -> NatOrd a -> NatOrd a -> NatOrd a forall a b. Coercible a b => a -> b coerce (forall a. Ord a => a -> a -> a max @a) {-# INLINE (\/) #-} newtype DivOrd a = DivOrd {forall a. DivOrd a -> a getDivOrd :: a} deriving (DivOrd a -> DivOrd a -> Bool (DivOrd a -> DivOrd a -> Bool) -> (DivOrd a -> DivOrd a -> Bool) -> Eq (DivOrd a) forall a. Eq a => DivOrd a -> DivOrd a -> Bool forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a $c== :: forall a. Eq a => DivOrd a -> DivOrd a -> Bool == :: DivOrd a -> DivOrd a -> Bool $c/= :: forall a. Eq a => DivOrd a -> DivOrd a -> Bool /= :: DivOrd a -> DivOrd a -> Bool Eq, Int -> DivOrd a -> ShowS [DivOrd a] -> ShowS DivOrd a -> String (Int -> DivOrd a -> ShowS) -> (DivOrd a -> String) -> ([DivOrd a] -> ShowS) -> Show (DivOrd a) forall a. Show a => Int -> DivOrd a -> ShowS forall a. Show a => [DivOrd a] -> ShowS forall a. Show a => DivOrd a -> String forall a. (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a $cshowsPrec :: forall a. Show a => Int -> DivOrd a -> ShowS showsPrec :: Int -> DivOrd a -> ShowS $cshow :: forall a. Show a => DivOrd a -> String show :: DivOrd a -> String $cshowList :: forall a. Show a => [DivOrd a] -> ShowS showList :: [DivOrd a] -> ShowS Show) instance (Integral a) => Poset (DivOrd a) where .<. :: DivOrd a -> DivOrd a -> Bool (.<.) (DivOrd a x) (DivOrd a y) = a -> a -> a forall a. Integral a => a -> a -> a rem a y a x a -> a -> Bool forall a. Eq a => a -> a -> Bool == a 0 {-# INLINE (.<.) #-} moebius :: forall i. Integral i => DivOrd a -> DivOrd a -> i moebius (DivOrd a x) (DivOrd a y) | Bool -> Bool not (Bool -> Bool) -> Bool -> Bool forall a b. (a -> b) -> a -> b $ a -> DivOrd a forall a. a -> DivOrd a DivOrd a x DivOrd a -> DivOrd a -> Bool forall a. Poset a => a -> a -> Bool .<. a -> DivOrd a forall a. a -> DivOrd a DivOrd a y = i 0 | Bool otherwise = [i] -> i forall a. Num a => [a] -> a forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a product ([i] -> i) -> ([a] -> [i]) -> [a] -> i forall b c a. (b -> c) -> (a -> b) -> a -> c . ([a] -> i) -> [[a]] -> [i] forall a b. (a -> b) -> [a] -> [b] map [a] -> i forall {a} {a}. Num a => [a] -> a mu ([[a]] -> [i]) -> ([a] -> [[a]]) -> [a] -> [i] forall b c a. (b -> c) -> (a -> b) -> a -> c . [a] -> [[a]] forall a. Eq a => [a] -> [[a]] L.group ([a] -> i) -> [a] -> i forall a b. (a -> b) -> a -> b $ a -> [a] forall i. Integral i => i -> [i] primeFactors (a -> a -> a forall a. Integral a => a -> a -> a quot a y a x) where mu :: [a] -> a mu [a _] = -a 1 mu [a] _ = a 0 instance (Integral a) => Lattice (DivOrd a) where /\ :: DivOrd a -> DivOrd a -> DivOrd a (/\) = (a -> a -> a) -> DivOrd a -> DivOrd a -> DivOrd a forall a b. Coercible a b => a -> b coerce (forall a. Integral a => a -> a -> a gcd @a) {-# INLINE (/\) #-} \/ :: DivOrd a -> DivOrd a -> DivOrd a (\/) = (a -> a -> a) -> DivOrd a -> DivOrd a -> DivOrd a forall a b. Coercible a b => a -> b coerce (forall a. Integral a => a -> a -> a lcm @a) {-# INLINE (\/) #-} instance FastZetaMoebius DivOrd where type Dim DivOrd = U.Vector Int fastZeta :: forall a (m :: * -> *). (Num a, Unbox a, PrimMonad m) => (Int -> DivOrd Int) -> Dim DivOrd -> MVector (PrimState m) a -> m () fastZeta Int -> DivOrd Int _ Dim DivOrd primes MVector (PrimState m) a g = do let n :: Int n = MVector (PrimState m) a -> Int forall a s. Unbox a => MVector s a -> Int UM.length MVector (PrimState m) a g Bool -> m () -> m () forall (f :: * -> *). Applicative f => Bool -> f () -> f () when (Int n Int -> Int -> Bool forall a. Ord a => a -> a -> Bool > Int 0) (m () -> m ()) -> m () -> m () forall a b. (a -> b) -> a -> b $ do a g0 <- MVector (PrimState m) a -> Int -> m a forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.read MVector (PrimState m) a g Int 0 Vector Int -> (Int -> m ()) -> m () forall (m :: * -> *) a b. (Monad m, Unbox a) => Vector a -> (a -> m b) -> m () U.forM_ Vector Int Dim DivOrd primes ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m () forall a b. (a -> b) -> a -> b $ \Int p -> Int -> (Int -> m ()) -> m () forall (m :: * -> *). Monad m => Int -> (Int -> m ()) -> m () rev (Int -> Int -> Int forall a. Integral a => a -> a -> a quot (Int n Int -> Int -> Int forall a. Num a => a -> a -> a - Int 1) Int p Int -> Int -> Int forall a. Num a => a -> a -> a + Int 1) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m () forall a b. (a -> b) -> a -> b $ \Int i -> do a c <- MVector (PrimState m) a -> Int -> m a forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) a g (Int p Int -> Int -> Int forall a. Num a => a -> a -> a * Int i) MVector (PrimState m) a -> (a -> a) -> Int -> m () forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> (a -> a) -> Int -> m () UM.unsafeModify MVector (PrimState m) a g (a -> a -> a forall a. Num a => a -> a -> a + a c) Int i MVector (PrimState m) a -> Int -> a -> m () forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> a -> m () UM.write MVector (PrimState m) a g Int 0 a g0 {-# INLINE fastZeta #-} fastMoebius :: forall a (m :: * -> *). (Num a, Unbox a, PrimMonad m) => (Int -> DivOrd Int) -> Dim DivOrd -> MVector (PrimState m) a -> m () fastMoebius Int -> DivOrd Int _ Dim DivOrd primes MVector (PrimState m) a f = do let n :: Int n = MVector (PrimState m) a -> Int forall a s. Unbox a => MVector s a -> Int UM.length MVector (PrimState m) a f Bool -> m () -> m () forall (f :: * -> *). Applicative f => Bool -> f () -> f () when (Int n Int -> Int -> Bool forall a. Ord a => a -> a -> Bool > Int 0) (m () -> m ()) -> m () -> m () forall a b. (a -> b) -> a -> b $ do a f0 <- MVector (PrimState m) a -> Int -> m a forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.read MVector (PrimState m) a f Int 0 Vector Int -> (Int -> m ()) -> m () forall (m :: * -> *) a b. (Monad m, Unbox a) => Vector a -> (a -> m b) -> m () U.forM_ Vector Int Dim DivOrd primes ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m () forall a b. (a -> b) -> a -> b $ \Int p -> Int -> (Int -> m ()) -> m () forall (m :: * -> *). Monad m => Int -> (Int -> m ()) -> m () rep (Int -> Int -> Int forall a. Integral a => a -> a -> a quot (Int n Int -> Int -> Int forall a. Num a => a -> a -> a - Int 1) Int p Int -> Int -> Int forall a. Num a => a -> a -> a + Int 1) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m () forall a b. (a -> b) -> a -> b $ \Int i -> do a c <- MVector (PrimState m) a -> Int -> m a forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) a f (Int p Int -> Int -> Int forall a. Num a => a -> a -> a * Int i) MVector (PrimState m) a -> (a -> a) -> Int -> m () forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> (a -> a) -> Int -> m () UM.unsafeModify MVector (PrimState m) a f (a -> a -> a forall a. Num a => a -> a -> a subtract a c) Int i MVector (PrimState m) a -> Int -> a -> m () forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> a -> m () UM.write MVector (PrimState m) a f Int 0 a f0 {-# INLINE fastMoebius #-} newtype SetOrd a = SetOrd {forall a. SetOrd a -> a getSetOrd :: a} deriving (SetOrd a -> SetOrd a -> Bool (SetOrd a -> SetOrd a -> Bool) -> (SetOrd a -> SetOrd a -> Bool) -> Eq (SetOrd a) forall a. Eq a => SetOrd a -> SetOrd a -> Bool forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a $c== :: forall a. Eq a => SetOrd a -> SetOrd a -> Bool == :: SetOrd a -> SetOrd a -> Bool $c/= :: forall a. Eq a => SetOrd a -> SetOrd a -> Bool /= :: SetOrd a -> SetOrd a -> Bool Eq, Int -> SetOrd a -> ShowS [SetOrd a] -> ShowS SetOrd a -> String (Int -> SetOrd a -> ShowS) -> (SetOrd a -> String) -> ([SetOrd a] -> ShowS) -> Show (SetOrd a) forall a. Show a => Int -> SetOrd a -> ShowS forall a. Show a => [SetOrd a] -> ShowS forall a. Show a => SetOrd a -> String forall a. (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a $cshowsPrec :: forall a. Show a => Int -> SetOrd a -> ShowS showsPrec :: Int -> SetOrd a -> ShowS $cshow :: forall a. Show a => SetOrd a -> String show :: SetOrd a -> String $cshowList :: forall a. Show a => [SetOrd a] -> ShowS showList :: [SetOrd a] -> ShowS Show) instance (Bits a) => Poset (SetOrd a) where .<. :: SetOrd a -> SetOrd a -> Bool (.<.) (SetOrd a x) (SetOrd a y) = a x a -> a -> a forall a. Bits a => a -> a -> a .&. a y a -> a -> Bool forall a. Eq a => a -> a -> Bool == a x {-# INLINE (.<.) #-} moebius :: forall i. Integral i => SetOrd a -> SetOrd a -> i moebius (SetOrd a x) (SetOrd a y) | Bool -> Bool not (Bool -> Bool) -> Bool -> Bool forall a b. (a -> b) -> a -> b $ a -> SetOrd a forall a. a -> SetOrd a SetOrd a x SetOrd a -> SetOrd a -> Bool forall a. Poset a => a -> a -> Bool .<. a -> SetOrd a forall a. a -> SetOrd a SetOrd a y = i 0 | Int -> Int -> Bool forall a. Bits a => a -> Int -> Bool testBit (a -> Int forall a. Bits a => a -> Int popCount (a -> Int) -> a -> Int forall a b. (a -> b) -> a -> b $ a -> a forall a. Bits a => a -> a complement a x a -> a -> a forall a. Bits a => a -> a -> a .&. a y) Int 0 = -i 1 | Bool otherwise = i 1 instance (Bits a) => Lattice (SetOrd a) where /\ :: SetOrd a -> SetOrd a -> SetOrd a (/\) = (a -> a -> a) -> SetOrd a -> SetOrd a -> SetOrd a forall a b. Coercible a b => a -> b coerce (forall a. Bits a => a -> a -> a (.&.) @a) {-# INLINE (/\) #-} \/ :: SetOrd a -> SetOrd a -> SetOrd a (\/) = (a -> a -> a) -> SetOrd a -> SetOrd a -> SetOrd a forall a b. Coercible a b => a -> b coerce (forall a. Bits a => a -> a -> a (.|.) @a) {-# INLINE (\/) #-} instance FastZetaMoebius SetOrd where type Dim SetOrd = Int fastZeta :: forall a (m :: * -> *). (Num a, Unbox a, PrimMonad m) => (Int -> SetOrd Int) -> Dim SetOrd -> MVector (PrimState m) a -> m () fastZeta Int -> SetOrd Int _ Dim SetOrd n MVector (PrimState m) a g = do Int -> (Int -> m ()) -> m () forall (m :: * -> *). Monad m => Int -> (Int -> m ()) -> m () rep Int Dim SetOrd n ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m () forall a b. (a -> b) -> a -> b $ \Int i -> do Int -> (Int -> m ()) -> m () forall (m :: * -> *). Monad m => Int -> (Int -> m ()) -> m () rep (Int -> Int -> Int forall a. Bits a => a -> Int -> a unsafeShiftL Int 1 Int Dim SetOrd n) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m () forall a b. (a -> b) -> a -> b $ \Int j -> do Bool -> m () -> m () forall (f :: * -> *). Applicative f => Bool -> f () -> f () unless (Int -> Int -> Bool forall a. Bits a => a -> Int -> Bool testBit Int j Int i) (m () -> m ()) -> m () -> m () forall a b. (a -> b) -> a -> b $ do a c <- MVector (PrimState m) a -> Int -> m a forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) a g (Int -> Int -> Int forall a. Bits a => a -> Int -> a setBit Int j Int i) MVector (PrimState m) a -> (a -> a) -> Int -> m () forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> (a -> a) -> Int -> m () UM.unsafeModify MVector (PrimState m) a g (a -> a -> a forall a. Num a => a -> a -> a + a c) Int j {-# INLINE fastZeta #-} fastMoebius :: forall a (m :: * -> *). (Num a, Unbox a, PrimMonad m) => (Int -> SetOrd Int) -> Dim SetOrd -> MVector (PrimState m) a -> m () fastMoebius Int -> SetOrd Int _ Dim SetOrd n MVector (PrimState m) a f = do Int -> (Int -> m ()) -> m () forall (m :: * -> *). Monad m => Int -> (Int -> m ()) -> m () rep Int Dim SetOrd n ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m () forall a b. (a -> b) -> a -> b $ \Int i -> do Int -> (Int -> m ()) -> m () forall (m :: * -> *). Monad m => Int -> (Int -> m ()) -> m () rep (Int -> Int -> Int forall a. Bits a => a -> Int -> a unsafeShiftL Int 1 Int Dim SetOrd n) ((Int -> m ()) -> m ()) -> (Int -> m ()) -> m () forall a b. (a -> b) -> a -> b $ \Int j -> do Bool -> m () -> m () forall (f :: * -> *). Applicative f => Bool -> f () -> f () unless (Int -> Int -> Bool forall a. Bits a => a -> Int -> Bool testBit Int j Int i) (m () -> m ()) -> m () -> m () forall a b. (a -> b) -> a -> b $ do a c <- MVector (PrimState m) a -> Int -> m a forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> Int -> m a UM.unsafeRead MVector (PrimState m) a f (Int -> Int -> Int forall a. Bits a => a -> Int -> a setBit Int j Int i) MVector (PrimState m) a -> (a -> a) -> Int -> m () forall (m :: * -> *) a. (PrimMonad m, Unbox a) => MVector (PrimState m) a -> (a -> a) -> Int -> m () UM.unsafeModify MVector (PrimState m) a f (a -> a -> a forall a. Num a => a -> a -> a subtract a c) Int j {-# INLINE fastMoebius #-}