{-|
module      :  Data.Number.Flint.Groups.DLog.FFI
copyright   :  (c) 2022 Hartmut Monien
license     :  GNU GPL, version 2 or above (see LICENSE)
maintainer  :  hmonien@uni-bonn.de
-}
module Data.Number.Flint.Groups.DLog.FFI (
  -- * Discrete logarithms mod ulong primes
  -- * Single evaluation
    dlog_once
  -- * Precomputations
  , DLogPrecomp (..)
  , CDLogPrecomp (..)
  , newDLogPrecompN
  , newDLogPrecompModpe
  , newDLogPrecompP
  , newDLogPrecompPE
  , newDLogPrecompSmall
  , withDLogPrecomp
  , dlog_precomp_n_init
  , dlog_precomp
  , dlog_precomp_clear
  , dlog_precomp_modpe_init
  , dlog_precomp_p_init
  , dlog_precomp_pe_init
  , dlog_precomp_small_init
  -- * Vector evaluations
  , dlog_vec_fill
  , dlog_vec_set_not_found
  , dlog_vec
  , dlog_vec_add
  , dlog_vec_loop
  , dlog_vec_loop_add
  , dlog_vec_eratos
  , dlog_vec_eratos_add
  , dlog_vec_sieve_add
  , dlog_vec_sieve
) where

-- Discrete logarithms mod ulong primes ----------------------------------------

import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.C.Types
import Foreign.Storable

import Data.Number.Flint.Flint
import Data.Number.Flint.NMod

#define DLOG_INLINES_C
#include <flint/dlog.h>

-- Single evaluation -----------------------------------------------------------

-- | /dlog_once/ /b/ /a/ /mod/ /n/ 
--
-- Return \(x\) such that \(b = a^x\) in
-- \((\mathbb Z/mod \mathbb Z)^\times\), where /a/ is known to have order
-- /n/.
foreign import ccall "dlog.h dlog_once"
  dlog_once :: CULong -> CULong -> Ptr CNMod -> CULong -> IO CULong

-- Precomputations -------------------------------------------------------------

data DLogPrecomp = DLogPrecomp {-# UNPACK #-} !(ForeignPtr CDLogPrecomp)
type CDLogPrecomp = CFlint DLogPrecomp

instance Storable CDLogPrecomp where
  sizeOf    _ = #{size      dlog_precomp_t}
  alignment _ = #{alignment dlog_precomp_t}
  peek = undefined
  poke = undefined
  
newDLogPrecompN a mod n num = do
  x <- mallocForeignPtr
  withForeignPtr x $ \x -> do
    dlog_precomp_n_init x a mod n num
  addForeignPtrFinalizer p_dlog_precomp_clear x
  return $ DLogPrecomp x

newDLogPrecompModpe a p e pe num = do
  x <- mallocForeignPtr
  withForeignPtr x $ \x -> do
    dlog_precomp_modpe_init x a p e pe num
  addForeignPtrFinalizer p_dlog_precomp_clear x
  return $ DLogPrecomp x

newDLogPrecompP a mod n num = do
  x <- mallocForeignPtr
  withForeignPtr x $ \x -> do
    dlog_precomp_p_init x a mod n num
  addForeignPtrFinalizer p_dlog_precomp_clear x
  return $ DLogPrecomp x

newDLogPrecompPE a mod p e pe num = do
  x <- mallocForeignPtr
  withForeignPtr x $ \x -> do
    dlog_precomp_pe_init x a mod p e pe num
  addForeignPtrFinalizer p_dlog_precomp_clear x
  return $ DLogPrecomp x

newDLogPrecompSmall a mod n num = do
  x <- mallocForeignPtr
  withForeignPtr x $ \x -> do
    dlog_precomp_small_init x a mod n num
  addForeignPtrFinalizer p_dlog_precomp_clear x
  return $ DLogPrecomp x

{-# INLINE withDLogPrecomp #-}
withDLogPrecomp (DLogPrecomp x) f = do
  withForeignPtr x $ \px -> f px >>= return . (DLogPrecomp x,)

--------------------------------------------------------------------------------

-- | /dlog_precomp_n_init/ /pre/ /a/ /mod/ /n/ /num/ 
--
-- Precompute data for /num/ discrete logarithms evaluations in the
-- subgroup generated by /a/ modulo /mod/, where /a/ is known to have order
-- /n/.
foreign import ccall "dlog.h dlog_precomp_n_init"
  dlog_precomp_n_init :: Ptr CDLogPrecomp
                      -> CULong -> CULong
                      -> CULong
                      -> CULong
                      -> IO ()

-- | /dlog_precomp/ /pre/ /b/ 
--
-- Return \(\log(b)\) for the group described in /pre/.
foreign import ccall "dlog.h dlog_precomp"
  dlog_precomp :: Ptr CDLogPrecomp -> CULong -> IO CULong

-- | /dlog_precomp_clear/ /pre/ 
--
-- Clears /t/.
foreign import ccall "dlog.h dlog_precomp_clear"
  dlog_precomp_clear :: Ptr CDLogPrecomp -> IO ()

foreign import ccall "dlog.h &dlog_precomp_clear"
  p_dlog_precomp_clear :: FunPtr (Ptr CDLogPrecomp -> IO ())

-- Specialized versions of @dlog_precomp_n_init@ are available when
-- specific information is known about the group:
--
-- | /dlog_precomp_modpe_init/ /pre/ /a/ /p/ /e/ /pe/ /num/ 
--
-- Assume that /a/ generates the group of residues modulo /pe/ equal
-- \(p^e\) for prime /p/.
foreign import ccall "dlog.h dlog_precomp_modpe_init"
  dlog_precomp_modpe_init :: Ptr CDLogPrecomp
                          -> CULong
                          -> CULong
                          -> CULong
                          -> CULong
                          -> CULong
                          -> IO ()

-- | /dlog_precomp_p_init/ /pre/ /a/ /mod/ /p/ /num/ 
--
-- Assume that /a/ has prime order /p/.
foreign import ccall "dlog.h dlog_precomp_p_init"
  dlog_precomp_p_init :: Ptr CDLogPrecomp
                      -> CULong
                      -> CULong
                      -> CULong
                      -> CULong
                      -> IO ()

-- | /dlog_precomp_pe_init/ /pre/ /a/ /mod/ /p/ /e/ /pe/ /num/ 
--
-- Assume that /a/ has primepower order /pe/ \(p^e\).
foreign import ccall "dlog.h dlog_precomp_pe_init"
  dlog_precomp_pe_init :: Ptr CDLogPrecomp
                       -> CULong
                       -> CULong
                       -> CULong
                       -> CULong
                       -> CULong
                       -> CULong
                       -> IO ()

-- | /dlog_precomp_small_init/ /pre/ /a/ /mod/ /n/ /num/ 
--
-- Make a complete lookup table of size /n/. If /mod/ is small, this is
-- done using an element-indexed array (see @dlog_table_t@), otherwise with
-- a sorted array allowing binary search.
foreign import ccall "dlog.h dlog_precomp_small_init"
  dlog_precomp_small_init :: Ptr CDLogPrecomp
                          -> CULong
                          -> CULong
                          -> CULong
                          -> CULong
                          -> IO ()

-- Vector evaluations ----------------------------------------------------------

-- These functions compute all logarithms of successive integers
-- \(1\dots n\).
--
-- | /dlog_vec_fill/ /v/ /nv/ /x/ 
--
-- Sets values /v[k]/ to /x/ for all /k/ less than /nv/.
foreign import ccall "dlog.h dlog_vec_fill"
  dlog_vec_fill :: Ptr CULong -> CULong -> CULong -> IO ()

-- | /dlog_vec_set_not_found/ /v/ /nv/ /mod/ 
--
-- Sets values /v[k]/ to @DLOG_NONE@ for all /k/ not coprime to /mod/.
foreign import ccall "dlog.h dlog_vec_set_not_found"
  dlog_vec_set_not_found :: Ptr CULong -> CULong -> Ptr CNMod -> IO ()

-- | /dlog_vec/ /v/ /nv/ /a/ /va/ /mod/ /na/ /order/ 
--
-- Sets /v[k]/ to \(\log(k,a)\) times value /va/ for \(0\leq k < nv\),
-- where /a/ has order /na/. /va/ should be /1/ for usual log computation.
foreign import ccall "dlog.h dlog_vec"
  dlog_vec :: Ptr CULong -> CULong -> CULong -> CULong -> Ptr CNMod -> CULong -> Ptr CNMod -> IO ()

-- | /dlog_vec_add/ /v/ /nv/ /a/ /va/ /mod/ /na/ /order/ 
--
-- Same parameters as before, but adds \(\log(k,a)\times v_a\) to /v[k]/
-- and reduce modulo /order/ instead of replacing the value. Indices /k/
-- such that /v[k]/ equals /DLOG_NONE/ are ignored.
foreign import ccall "dlog.h dlog_vec_add"
  dlog_vec_add :: Ptr CULong -> CULong -> CULong -> CULong -> Ptr CNMod -> CULong -> Ptr CNMod -> IO ()

-- Depending on the relative size of /nv/ and /na/, these two /dlog_vec/
-- functions call one of the following functions.
--
-- | /dlog_vec_loop/ /v/ /nv/ /a/ /va/ /mod/ /na/ /order/ 
--
foreign import ccall "dlog.h dlog_vec_loop"
  dlog_vec_loop :: Ptr CULong -> CULong -> CULong -> CULong -> Ptr CNMod -> CULong -> Ptr CNMod -> IO ()

-- | /dlog_vec_loop_add/ /v/ /nv/ /a/ /va/ /mod/ /na/ /order/ 
--
-- Perform a complete loop of size /na/ on powers of /a/ to fill the
-- logarithm values, discarding powers outside the bounds of /v/. This
-- requires no discrete logarithm computation.
foreign import ccall "dlog.h dlog_vec_loop_add"
  dlog_vec_loop_add :: Ptr CULong -> CULong -> CULong -> CULong -> Ptr CNMod -> CULong -> Ptr CNMod -> IO ()

-- | /dlog_vec_eratos/ /v/ /nv/ /a/ /va/ /mod/ /na/ /order/ 
--
foreign import ccall "dlog.h dlog_vec_eratos"
  dlog_vec_eratos :: Ptr CULong -> CULong -> CULong -> CULong -> Ptr CNMod -> CULong -> Ptr CNMod -> IO ()

-- | /dlog_vec_eratos_add/ /v/ /nv/ /a/ /va/ /mod/ /na/ /order/ 
--
-- Compute discrete logarithms of prime numbers less than /nv/ and
-- propagate to composite numbers.
foreign import ccall "dlog.h dlog_vec_eratos_add"
  dlog_vec_eratos_add :: Ptr CULong -> CULong -> CULong -> CULong -> Ptr CNMod -> CULong -> Ptr CNMod -> IO ()

-- | /dlog_vec_sieve_add/ /v/ /nv/ /a/ /va/ /mod/ /na/ /order/ 
--
foreign import ccall "dlog.h dlog_vec_sieve_add"
  dlog_vec_sieve_add :: Ptr CULong -> CULong -> CULong -> CULong -> Ptr CNMod -> CULong -> Ptr CNMod -> IO ()

-- | /dlog_vec_sieve/ /v/ /nv/ /a/ /va/ /mod/ /na/ /order/ 
--
-- Compute the discrete logarithms of the first few prime numbers, then use
-- them as a factor base to obtain the logarithms of larger primes by
-- sieving techniques.
-- 
-- In the the present implementation, the full index-calculus method is not
-- implemented.
foreign import ccall "dlog.h dlog_vec_sieve"
  dlog_vec_sieve :: Ptr CULong -> CULong -> CULong -> CULong -> Ptr CNMod -> CULong -> Ptr CNMod -> IO ()