MemoryOddity.hs

{-
Requires (from Stack lts-12.11)
- arithmoi >= 0 && < 1
- containers >= 0.5 && < 0.6
- parallel >= 3 && < 4
- split >= 0.2 && < 0.3
-}
import qualified Math.NumberTheory.Primes.Sieve as Math

import Data.IntSet (IntSet)
import qualified Data.IntSet as Set
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as Map

import Data.List (unfoldr)
import Data.List.Split

import Control.Parallel.Strategies
import GHC.Conc (numCapabilities)

data Interval = Interval { start :: !Int, end :: !Int }
  deriving (Show, Eq, Ord)

inside :: Int -> Interval -> Bool
inside p (Interval s e) = p >= s && p <= e

merge :: Interval -> Interval -> Interval
merge i1 i2 | i1 <= i2 = Interval (start i1) (end i2)
            | otherwise = Interval (start i2) (end i1)

contains :: Interval -> Interval -> Bool
contains big small = start big <= start small && end big >= end small

-- Partition Tree, similar to a segment tree, but intervals on leaves
-- are all disjoint. Furthermore, this makes the finger tree part of the structure a bit useless
-- So the tags are only stored on the leaves
data PTree
  = Leaf { tag :: !Int, interval :: !Interval }
  | Branch { interval :: !Interval , left :: !PTree , right :: !PTree }
  deriving (Show)

union :: PTree -> PTree -> PTree
union l r = Branch (merge (interval l) (interval r)) l r

-- Assumes that n is a part of the partition
query :: Int -> PTree -> Int
query _ (Leaf t _) = t
query n (Branch _ l r) | n <= (end $ interval l) = query n l
                       | otherwise = query n r

fromList :: [(Int, Interval)] -> PTree
fromList is = head $ head $ dropWhile (not . converged) $ iterate (unfoldr connect) $ fmap (uncurry Leaf) is
  where
    connect :: [PTree] -> Maybe (PTree, [PTree])
    connect [] = Nothing
    connect [x,y,z] = Just (x `union` y `union` z, [])
    connect (x:y:rest) = Just (x `union` y, rest)

    converged :: [a] -> Bool
    converged [_] = True
    converged _ = False


p609 :: Int -> Int -> Int
p609 n m =
  let dicts = withStrategy (parList rdeepseq) $ foldr c Map.empty <$> chunksOf (n `div` numCapabilities) [1..n]
  in Map.foldr (\x prod -> ((x `mod` m) * prod) `mod` m) 1 $ Map.unionsWith (+) $ dicts
  where
    primes :: [Int]
    primes = fromIntegral <$> Math.primes

    primeSet :: IntSet
    primeSet = Set.fromList $ take n primes

    p :: Int -> Int
    p n = if Set.member n primeSet then 0 else 1

    intervals :: ([Int], Int, Int) -> Maybe ((Int, Interval), ([Int], Int, Int))
    intervals (_, s, c) | s > n = Nothing
    intervals (ps, s, c) = 
      let (ps', e) = go ps s
      in Just ((c, Interval s (e - 1)), (ps', e, c + 1))
      where
        go (p:ps) i | i < p = go (p:ps) (i+1)
                    | i >= p || i == n = (ps, i)
                    
    partition :: PTree
    partition = fromList $ unfoldr intervals (primes, 1, 0)

    c :: Int -> IntMap Int -> IntMap Int
    c u_i table = go (query u_i partition) table (p u_i) 
      where
        go 0 table _ = table
        go u_i table count =
          let count' = count + (p u_i)
          in go (query u_i partition) (Map.insertWith (+) count' 1 table) count'