<commit_before>{-# LANGUAGE BangPatterns, RankNTypes, GeneralizedNewtypeDeriving, ScopedTypeVariables #-}
{-# LANGUAGE TupleSections, DeriveFunctor #-}
module ECC.Code.LDPC.Fast.CachedMult where

-- Uses the StableDiv data structure to cache multiplications.

import ECC.Code.LDPC.Utils
import ECC.Types
import ECC.Puncture
import Data.Char (isDigit)
import Data.Matrix
import Data.Bit
import Data.Bits
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector as V
import qualified Data.Matrix.QuasiCyclic as Q
import Debug.Trace

import qualified ECC.Code.LDPC.Fast.Encoder as E

import Data.Semigroup
import Data.Foldable (foldl')


data StableDiv
  = StableDiv
       !Double -- | The "worst" value for division: the closest to zero
       !Double -- | The result of the multiplications,
               --   excluding the "worst" value

absMinMax :: Double -> Double -> (Double, Double)
absMinMax x y
  | abs x < abs y = (x, y)
  | otherwise     = (y, x)

instance Semigroup StableDiv where
  StableDiv a b <> StableDiv c d =
    StableDiv minOfMins (b * maxOfMins * d)
    where
      (minOfMins, maxOfMins) = absMinMax a c

lit :: Double -> StableDiv
lit x
  | x >= 1    = StableDiv 1 x
  | otherwise = StableDiv x 1

sdiv :: StableDiv -> Double -> Double
sdiv (StableDiv a b) c
  | a == c    = b
  | otherwise = a * (b/c)



type M a = Matrix a
type V a = U.Vector a

-- An arraylet is a small array, that acts like a NxN matrix, with one element in each row and column,
-- (a rotated identity matrix, for example). We index explicity start at row/column *0*.
-- The Int argument is the rotation. The row index is the index into the array.
data Arraylet a = Arraylet Int (U.Vector a)
 deriving (Show)

mapArraylet :: (U.Unbox a, U.Unbox b) => (a -> b) -> Arraylet a -> Arraylet b
mapArraylet f (Arraylet sz a) = Arraylet sz (U.map f a)

sizeOfArraylet :: Arraylet a -> Int
sizeOfArraylet (Arraylet n _) = n

lookupArraylet :: U.Unbox a => Arraylet a -> (Int,Int) -> Maybe a
lookupArraylet (Arraylet off m) (r,c)
     | r < 0 || c < 0           = error "arraylet lookup out of bounds"
     | r >= len || c >= len     = error "arraylet lookup out of bounds"
     | r == (c - off) `mod` len = Just $ m U.! r
     | otherwise                = Nothing
  where len = U.length m

arrayArraylet :: forall a . U.Unbox a => Int -> Int -> ((Int,Int) -> a) -> Arraylet a
arrayArraylet sz off k = Arraylet off $ U.generate sz $ \ r -> k (r,(r + off) `mod` sz)

foldRowsArraylet :: U.Unbox a => Arraylet a -> U.Vector a
foldRowsArraylet (Arraylet sz m) = b `mappend` a
    where (a,b) = U.splitAt (U.length m - sz) m

foldRowsArraylet' :: U.Unbox a => ((Int,Int) -> a -> b) -> Arraylet a -> V.Vector b
foldRowsArraylet' k (Arraylet off m) = V.generate sz $ \ ix -> let !ix' = (ix + (sz - off)) `mod` sz in
                  k (ix',(ix' + off) `mod` sz) (m U.! ix')
  where !sz = U.length m

foldColsArraylet :: U.Unbox a => Arraylet a -> U.Vector a
foldColsArraylet (Arraylet sz m) = m

foldColsArraylet' :: U.Unbox a => ((Int,Int) -> a -> b) -> Arraylet a -> V.Vector b
foldColsArraylet' k (Arraylet off m) = V.generate sz $ \ ix -> k (ix,(ix + off) `mod` sz) (m U.! ix)
  where !sz = U.length m

--indexByRow :: Arraylet a -> Int -> a
--indexByRow (Arraylet a _) = undefined

-- A matrix of arraylets, all the same size. 0-indexed.
data Matrixlet a = Matrixlet Int (Matrix (Maybe (Arraylet a)))
 deriving Show

initMatrixlet :: forall a . U.Unbox a => a -> Q.QuasiCyclic Integer -> Matrixlet a
initMatrixlet zero (Q.QuasiCyclic sz m) = Matrixlet sz (fmap f m)
    where f :: Integer -> Maybe (Arraylet a)
          f 0 = Nothing
          f n | popCount n == 1 = Just $ Arraylet (g n) (U.replicate sz zero)
              | otherwise       = error $ "QuasiCyclic matrix has non-powers of two initial value of " ++ show n

          -- The number must be a power of two, because there is only one bit set.
          g :: Integer -> Int
          g x | x `testBit` 0 = 0
              | x == 0        = error "got to zero; should never happen"
              | otherwise = 1 + g (x `shiftR` 1)


matrixMatrixlet :: forall a b . U.Unbox b => Matrixlet a -> ((Int,Int) -> b) -> Matrixlet b
matrixMatrixlet (Matrixlet sz a) k = Matrixlet sz $ matrix (nrows a) (ncols a) k'
  where
      k' :: (Int, Int) -> Maybe (Arraylet b)
      k' (r,c) = case a ! (r,c) of
                   Nothing -> Nothing
                   Just (Arraylet off _) -> Just $ arrayArraylet sz off $ \ (r',c') -> k ((r-1) * sz + r',(c-1) * sz + c')

mapMatrixlet :: (U.Unbox a, U.Unbox b) => (a -> b) -> Matrixlet a -> Matrixlet b
mapMatrixlet f (Matrixlet sz a) = Matrixlet sz (fmap (fmap (mapArraylet f)) a)

-- Assumes there is always one value on every column
foldColsMatrixlet :: U.Unbox a => (a -> a -> a) -> Matrixlet a -> U.Vector a
foldColsMatrixlet f (Matrixlet n a) = U.concat
    [ foldr1 (U.zipWith f) [ foldColsArraylet x | c <- [1..ncols a], Just x <- [a ! (r,c)]]
    | r <- [1..nrows a]
    ]

foldColsMatrixlet' :: U.Unbox a => ((Int,Int) -> a -> b) -> (b -> b -> b) -> Matrixlet a -> V.Vector b
foldColsMatrixlet' f g (Matrixlet sz a) = V.concat
    [ foldr1 (V.zipWith g) [ foldColsArraylet' (\ (r',c') a -> f ((r-1) * sz + r',(c-1) * sz + c') a) x | c <- [1..ncols a], Just x <- [a ! (r,c)]]
    | r <- [1..nrows a]
    ]
<commit_msg>Indexed maps over arraylets and matrixlets<commit_after>{-# LANGUAGE BangPatterns, RankNTypes, GeneralizedNewtypeDeriving, ScopedTypeVariables #-}
{-# LANGUAGE TupleSections, DeriveFunctor #-}
module ECC.Code.LDPC.Fast.CachedMult where

-- Uses the StableDiv data structure to cache multiplications.

import ECC.Code.LDPC.Utils
import ECC.Types
import ECC.Puncture
import Data.Char (isDigit)
import Data.Matrix
import Data.Bit
import Data.Bits
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector as V
import qualified Data.Matrix.QuasiCyclic as Q
import Debug.Trace

import qualified ECC.Code.LDPC.Fast.Encoder as E

import Data.Semigroup
import Data.Foldable (foldl')


data StableDiv
  = StableDiv
       !Double -- | The "worst" value for division: the closest to zero
       !Double -- | The result of the multiplications,
               --   excluding the "worst" value

absMinMax :: Double -> Double -> (Double, Double)
absMinMax x y
  | abs x < abs y = (x, y)
  | otherwise     = (y, x)

instance Semigroup StableDiv where
  StableDiv a b <> StableDiv c d =
    StableDiv minOfMins (b * maxOfMins * d)
    where
      (minOfMins, maxOfMins) = absMinMax a c

lit :: Double -> StableDiv
lit x
  | x >= 1    = StableDiv 1 x
  | otherwise = StableDiv x 1

sdiv :: StableDiv -> Double -> Double
sdiv (StableDiv a b) c
  | a == c    = b
  | otherwise = a * (b/c)



type M a = Matrix a
type V a = U.Vector a

-- An arraylet is a small array, that acts like a NxN matrix, with one element in each row and column,
-- (a rotated identity matrix, for example). We index explicity start at row/column *0*.
-- The Int argument is the rotation. The row index is the index into the array.
data Arraylet a = Arraylet Int (U.Vector a)
 deriving (Show)

mapArraylet :: (U.Unbox a, U.Unbox b) => (a -> b) -> Arraylet a -> Arraylet b
mapArraylet f (Arraylet sz a) = Arraylet sz (U.map f a)

sizeOfArraylet :: Arraylet a -> Int
sizeOfArraylet (Arraylet n _) = n

lookupArraylet :: U.Unbox a => Arraylet a -> (Int,Int) -> Maybe a
lookupArraylet (Arraylet off m) (r,c)
     | r < 0 || c < 0           = error "arraylet lookup out of bounds"
     | r >= len || c >= len     = error "arraylet lookup out of bounds"
     | r == (c - off) `mod` len = Just $ m U.! r
     | otherwise                = Nothing
  where len = U.length m

arrayArraylet :: forall a . U.Unbox a => Int -> Int -> ((Int,Int) -> a) -> Arraylet a
arrayArraylet sz off k = Arraylet off $ U.generate sz $ \ r -> k (r,(r + off) `mod` sz)

-- | An indexed map over arraylets (indexing based on 'arrayArraylet')
imapArraylet :: forall a b . (U.Unbox a, U.Unbox b) => Int -> ((Int,Int) -> a -> b) -> Arraylet a -> Arraylet b
imapArraylet sz k (Arraylet off arr) = Arraylet off $ U.imap (\ r v -> k (r,(r + off) `mod` sz) v) arr

foldRowsArraylet :: U.Unbox a => Arraylet a -> U.Vector a
foldRowsArraylet (Arraylet sz m) = b `mappend` a
    where (a,b) = U.splitAt (U.length m - sz) m

foldRowsArraylet' :: U.Unbox a => ((Int,Int) -> a -> b) -> Arraylet a -> V.Vector b
foldRowsArraylet' k (Arraylet off m) = V.generate sz $ \ ix -> let !ix' = (ix + (sz - off)) `mod` sz in
                  k (ix',(ix' + off) `mod` sz) (m U.! ix')
  where !sz = U.length m

foldColsArraylet :: U.Unbox a => Arraylet a -> U.Vector a
foldColsArraylet (Arraylet sz m) = m

foldColsArraylet' :: U.Unbox a => ((Int,Int) -> a -> b) -> Arraylet a -> V.Vector b
foldColsArraylet' k (Arraylet off m) = V.generate sz $ \ ix -> k (ix,(ix + off) `mod` sz) (m U.! ix)
  where !sz = U.length m

--indexByRow :: Arraylet a -> Int -> a
--indexByRow (Arraylet a _) = undefined

-- A matrix of arraylets, all the same size. 0-indexed.
data Matrixlet a = Matrixlet Int (Matrix (Maybe (Arraylet a)))
 deriving Show

initMatrixlet :: forall a . U.Unbox a => a -> Q.QuasiCyclic Integer -> Matrixlet a
initMatrixlet zero (Q.QuasiCyclic sz m) = Matrixlet sz (fmap f m)
    where f :: Integer -> Maybe (Arraylet a)
          f 0 = Nothing
          f n | popCount n == 1 = Just $ Arraylet (g n) (U.replicate sz zero)
              | otherwise       = error $ "QuasiCyclic matrix has non-powers of two initial value of " ++ show n

          -- The number must be a power of two, because there is only one bit set.
          g :: Integer -> Int
          g x | x `testBit` 0 = 0
              | x == 0        = error "got to zero; should never happen"
              | otherwise = 1 + g (x `shiftR` 1)


matrixMatrixlet :: forall a b . U.Unbox b => Matrixlet a -> ((Int,Int) -> b) -> Matrixlet b
matrixMatrixlet (Matrixlet sz a) k = Matrixlet sz $ matrix (nrows a) (ncols a) k'
  where
      k' :: (Int, Int) -> Maybe (Arraylet b)
      k' (r,c) = case a ! (r,c) of
                   Nothing -> Nothing
                   Just (Arraylet off _) -> Just $ arrayArraylet sz off $ \ (r',c') -> k ((r-1) * sz + r',(c-1) * sz + c')

-- | Indexed map over matrixlets (based on 'matrixMatrixlet')
imapMatrixlet :: forall a b . (U.Unbox a, U.Unbox b) => Matrixlet a -> ((Int,Int) -> a -> b) -> Matrixlet b
imapMatrixlet (Matrixlet sz a) k = Matrixlet sz $ matrix (nrows a) (ncols a) k'
  where
      k' :: (Int, Int) -> Maybe (Arraylet b)
      k' (r,c) =
        case a ! (r,c) of
          Nothing -> Nothing
          Just arr ->
            Just $ flip (imapArraylet sz) arr $ \ (r',c') v ->
              let r'' = (r-1) * sz + r'
                  c'' = (c-1) * sz + c'
              in
              k (r'',c'') v

mapMatrixlet :: (U.Unbox a, U.Unbox b) => (a -> b) -> Matrixlet a -> Matrixlet b
mapMatrixlet f (Matrixlet sz a) = Matrixlet sz (fmap (fmap (mapArraylet f)) a)

-- Assumes there is always one value on every column
foldColsMatrixlet :: U.Unbox a => (a -> a -> a) -> Matrixlet a -> U.Vector a
foldColsMatrixlet f (Matrixlet n a) = U.concat
    [ foldr1 (U.zipWith f) [ foldColsArraylet x | c <- [1..ncols a], Just x <- [a ! (r,c)]]
    | r <- [1..nrows a]
    ]

foldColsMatrixlet' :: U.Unbox a => ((Int,Int) -> a -> b) -> (b -> b -> b) -> Matrixlet a -> V.Vector b
foldColsMatrixlet' f g (Matrixlet sz a) = V.concat
    [ foldr1 (V.zipWith g) [ foldColsArraylet' (\ (r',c') a -> f ((r-1) * sz + r',(c-1) * sz + c') a) x | c <- [1..ncols a], Just x <- [a ! (r,c)]]
    | r <- [1..nrows a]
    ]

{'fraction_non_alphanumeric': 0.14143598615916955, 'fraction_numerical': 0.004325259515570935, 'mean_word_length': 2.5033333333333334, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 163, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5814538', 'n_tokens_mistral': 4149, 'n_tokens_neox': 4062, 'n_words': 1530}