{-# LANGUAGE Safe, ScopedTypeVariables, MagicHash, UnboxedTuples, TypeFamilies, FlexibleContexts, CPP #-}

-- | A memory manager operating inside memory-mapped files.
--
--   Notes on timing and locking
--
--   * The library synchronizes calls to 'malloc' and 'free', however the caller has to do any other synchronization itself.
--
--   * As usual file should be closed in order to see writes transferred to disk.
module File.Mapped (Represent, Pointer, FilePtr, pointer, offset, fileSrc, Mapped, path, openMapped, closeMapped, malloc, make, free,
-- * Pointer operations
fpeek, fpoke, peekPtr, pokePtr, coerce, plus, minus, toRepr,
-- * Signal masking
maskSigint) where

import Foreign.Ptr
import qualified Foreign.Storable as S
import System.IO
import System.IO.Error
import System.FilePath.Posix
import System.FileLock
import Control.Exception (try, finally, throwIO)
import Control.Monad
import Data.Word
import Data.Int
import Data.Bits
import Data.Char
import Data.Storables
import Data.MyEndianness
import File.Mapper
import qualified Data.ByteString as B

#ifndef __WIN32__
import System.Posix.Signals
#endif

-- The bit representation of a data type (what the bits look like on disk).
type family Represent a :: *
type instance Represent Int8 = Int8
type instance Represent Int16 = Int16
type instance Represent Int32 = Int32
type instance Represent Int = Int32
type instance Represent Int64 = Int64
type instance Represent Word8 = Word8
type instance Represent Word16 = Word16
type instance Represent Pointer = Pointer
type instance Represent Word64 = Word64
type instance Represent (FilePtr t) = Pointer
type instance Represent Float = Float
type instance Represent Char = Char
type instance Represent Bool = Bool
type instance Represent (t, u) = (Represent t, Represent u)
type instance Represent (t, u, v) = (Represent t, Represent u, Represent v)
type instance Represent (t, u, v, w) = (Represent t, Represent u, Represent v, Represent w)
type instance Represent (t, u, v, w, x) = (Represent t, Represent u, Represent v, Represent w, Represent x)
type instance Represent (t, u, v, w, x, y) = (Represent t, Represent u, Represent v, Represent w, Represent x, Represent y)
type instance Represent (t, u, v, w, x, y, z) = (Represent t, Represent u, Represent v, Represent w, Represent x, Represent y, Represent z)
type instance Represent () = ()

-- | A 'FilePtr' keeps track of:
--
-- * The address of data where it is mapped into memory.
--
-- * The address (offset) of the data on disk.
--
-- * Information about the file being referenced.
data FilePtr a = FilePtr !(Ptr (Represent a)) !Pointer !Mapped

type Header = (Word64, Pointer, Pointer, Word32)

pointer ~(FilePtr p _ _) = p

offset ~(FilePtr _ off _) = off

fileSrc ~(FilePtr _ _ file) = file

instance Eq (FilePtr a) where
	FilePtr _ p mp == FilePtr _ p2 mp2 = p == p2 && path mp == path mp2

instance Show (FilePtr a) where
	showsPrec _ (FilePtr _ p file) = showsPrec 11 p . ("@"++) . showsPrec 11 (path file)

count pred f = fst . head . dropWhile (not . pred . snd) . zip [0..] . iterate f

onePos :: Word64 -> Int
onePos x | x >= 0 = count (`testBit` 0) (`shiftR` 1) x

nBits :: (Integral i, Bits i) => i -> i
nBits x | x >= 0 = if popCount x == 1 then n else n + 1 where
	n = count (==1) (`shiftR` 1) x

-- Allocator playing the role of sbrk:
dumbAllocator mapped _ = liftM (\(p, off) -> FilePtr p (fromIntegral off) mapped) $ newBlock (table mapped)

-- Dumb freer that never frees.
dumbFreer _ = return False

data Mapped = Mapped !Table !(Ptr Pointer)

table ~(Mapped tbl _) = tbl

block ~(Mapped _ blk) = blk

path ~(Mapped (Table p _) _) = p

-- | Open a memory mapped file, returning the root pointer.
openMapped :: FilePath -> IO (FilePtr (B.ByteString, t))
openMapped path = do
	tab <- newTable path
	ei <- try (openBinaryFile path ReadMode)
	setup <- either
		(\ex -> if isDoesNotExistError ex then return Nothing else throwIO ex)
		(\hdl -> liftM (\n -> if n == 0 then Nothing else Just hdl) $ hFileSize hdl)
		ei
	maybe
		(maskSigint $ withFileLock (path ++ ".lock") Exclusive $ const $ do
			FilePtr p _ _ :: FilePtr Header <- dumbAllocator (Mapped tab nullPtr) (2 ^ 16)
			sequence_ $ zipWith (\i (n :: Word8) -> S.pokeByteOff (castPtr p) i n) [24..] (map (fromIntegral . ord) "mmap")

			let file = Mapped tab (castPtr p)

			-- Allocate an outer superblock with one block allocated.
			pokeByteOff (castPtr p) 64 (8 :: Pointer)
			poke p (complement 1, 64, 0, 10)

			-- Allocate a root object. This is an eight-byte object. By convention, the first DWORD
			-- is an identifying ByteString. The second DWORD is free for any use.
			_malloc (_malloc dumbAllocator initialPattern1 10 16) initialPattern 5 9 file 8)
		(\hdl -> do
			hClose hdl
			p <- mapBlock tab 0
			s :: [Word8] <- mapM (S.peekByteOff p) [24..27]
			unless (s == map (fromIntegral . ord) "mmap") $ error "File.Mapped.openMapped: wrong file type"
			return $ FilePtr (castPtr (plusPtr p off)) (fromIntegral off) (Mapped tab p))
		setup
 where
	off = 1024 + 24 + 24

-- | Close a memory mapped file. This will invalidate all pointers created on this memory map.
closeMapped :: Mapped -> IO ()
closeMapped (Mapped tab _) = unmapAll tab

-- Let f be the heap at the start of appendHdrList. Let Integrity hold at the outset.
appendHdrList :: Mapped -> Int -> Ptr Header -> Pointer -> Word64 -> IO ()
appendHdrList mapped n blk blkOff m = do
	-- Get the header pointer
	let hdrLoc = plusPtr (block mapped) (24 + 4 * n) -- pointer to next pointer
	hdr <- peek hdrLoc -- next pointer

	-- Update the bits
	poke hdrLoc (blkOff + 8) -- poke next pointer to previous pointer
	poke blk (m, fromIntegral $ minusPtr hdrLoc (block mapped), hdr, fromIntegral n) -- poke header
	unless (hdr == 0) $ do
		p <- mapInBlock mapped hdr
		poke p (blkOff + 12) -- poke prev pointer to next pointer
-- The heap at the end is f[f (6+n)/blk+2][blk/m][blk+1/m2][blk+2/6+n][blk+3/f (6 + n)][blk+4/n][f (6 + n)/blk+3].
--
-- appendHdrList is assumed to yield an allocated superblock at blk.
--
-- Demonstrations:
-- Integrity1: Case: blk+3. f (f (blk+3)) = f (f (6+n)) = f (blk+3) = f (6+n) = blk+3. Otherwise f is untouched.
-- Integrity2/3: Case: block allocated at blk. Sub-case: hdr is non-null. f (6 + n) - 2 is allocated by Integrity6/7. f (6 + n) = hdr. f (blk + 3) = hdr. Therefore f (blk + 3) - 2 is allocated. Otherwise f is untouched.
-- Integrity4: Case: 6 + n. Sub-case: hdr is non-null. f (f (6 + n) + 2) = n by Integrity8. f (6 + n) = f (blk + 3). n = f (blk + 4). Therefore f (f (blk + 3)  2) = f (blk + 4). Otherwise f is untouched.
-- Integrity6/7: Case: 6 + n. f (6 + n) - 2 = blk which is alloced by assumption. Otherwise f untouched.
-- Integrity8: Case: 6 + n. f (f (6 + n) + 2) = f (blk + 2 + 2) = f (blk + 4) = n. Otherwise f untouched.
-- AllocatedIntegrity: Unchanged.
-- End.

initialPattern n = case n of
	-- 3 -> 2 ^ 61 - 1
	-- 4 -> 2 ^ 62 - 1
	5 -> 2 ^ 31 - 1
	6 -> 32767
	7 -> 127
	8 -> 7
	9 -> 1
	_ -> 1

initialPattern1 n = case n of
	10 -> complement 0
	11 -> 2 ^ 32 - 1
	12 -> 65535
	13 -> 255
	14 -> 15
	15 -> 3
	_ -> 1

mapInBlock :: Mapped -> Pointer -> IO (Ptr Word32)
mapInBlock = mapBlock . table

fixPtrs mapped prev next = do
	p <- mapInBlock mapped prev
	poke p next
	unless (next == 0) $ do
		p <- mapInBlock mapped next
		poke p prev

-- Let f = the heap at the start. Let Integrity be valid at the outset. (Shallow) inductive hypothesis: let alloc be able to allocate at least as large as (2 ^ mx).
-- The main allocator:
_malloc alloc pattern mn mx mapped sz = if sz <= 0 || sz > 2 ^ 16 - 24 then
		error "File.Mapped.malloc: allocation out of range"
	else do
	-- Find a free superblock
	let bits = nBits (sz + 48) `max` mn
	blkOff <- peekByteOff (block mapped) (24 + 4 * bits)
	(p, off) <- if blkOff == 0 then do
			FilePtr p off _ :: FilePtr Header <- alloc mapped ((2 ^ mx - 24) `max` sz)
			appendHdrList mapped bits p off (pattern bits)
			return (p, off)
		else do
			p <- mapInBlock mapped (blkOff - 8)
			return (castPtr p, blkOff - 8)

	(m, prev, next, _) :: Header <- peek p

	-- Allocate a block
	let pos = onePos m
	let m' = clearBit m (fromIntegral pos)
	poke p (m', prev, next, fromIntegral bits)

	-- If this superblock is now full, fix up the pointers
	when (m' == 0) $ fixPtrs mapped prev next

	return (FilePtr (plusPtr p (fromIntegral (pos * 2 ^ bits + 24))) (off + fromIntegral (pos * 2 ^ bits + 24)) mapped)
-- The result is assumed to be allocated with size at least sz.
--
-- The heap at the end is f[*/m'][*+1/m2'][*+4/n] for some choice of n so that f (6 + n) is non-null, modulo call to appendHdrList and fixPtrs. In the case that a block is available, let * = f (6 + n) - 2 otherwise * = the result of alloc.
--
-- Demonstrations:
-- Integrity_n: Appeal to the correctness of appendHdrList, fixPtrs, and the inductive hypothesis.
-- AllocatedIntegrity: Only the current allocation and the next allocation over need to be considered.
-- The next allocation over, p2, should be >= p + sz, unless the current block is already allocated.
--
-- Case: the current block is already allocated.
--
-- I know the header being read is valid by Integrity. 'onePos' will find a one bit in the
-- m' vector, which is unallocated by definition. Therefore this case is refuted.
--
-- Case: the current block is unallocated.
--
-- I need to prove that p2 >= p + sz. The allocations are spaced 2 ^ n bytes apart. This is true because
-- n was set at least nBits (sz + 24). 'nBits' has the property that x <= 2 ^ nBits x.
-- Therefore p2 >= p + sz + 24.
-- End.

truncateOffset mx off = off `shiftR` mx `shiftL` mx

-- Let 'freer' be able to free the next superblock size up (10 bytes or 16 bytes)
-- [shallow inductive hypothesis]. "Freeing" is defined as making the
-- block unallocated. Let Integrity hold at the outset.
_free freer pattern truncate (FilePtr p off mapped) = do
	-- Figure out the beginning of the block
	let blk = plusPtr nullPtr $ truncate $ minusPtr p nullPtr
	let blkOff = fromIntegral $ truncate $ fromIntegral off

	(m, prev, next, n) :: Header <- peek blk

	let m' = setBit m (fromIntegral $ (off - blkOff - 24) `shiftR` fromIntegral n)
	poke blk (m', prev, next, n)
	if m' == pattern n then do
			-- If the superblock is completely empty, try to free it
			b <- freer (FilePtr blk blkOff mapped :: FilePtr Header)
			when b $ fixPtrs mapped prev next
		else
			-- Add the superblock to the header list
			when (m == 0) $ appendHdrList mapped (fromIntegral n) blk blkOff m'

	return True
-- Integrity (provided the passed pointer is allocated):
-- 'truncate' is assumed to locate the header in which this block is located.
-- Having located the superblock, appeal to the inductive hypothesis and Integrity
-- for fixPtrs and appendHdrList.
-- End.

-- | Make a FilePtr from a Pointer.
make :: Mapped -> Pointer -> IO (FilePtr t)
make mapped off = liftM (\p -> FilePtr (castPtr p) off mapped) $ mapInBlock mapped off

-- The allocator is obtained by nesting twice.
--
-- | Allocate a block.
malloc mapped n = maskSigint $ withFileLock (path mapped ++ ".lock") Exclusive $ const $ do
	p <- _malloc (_malloc dumbAllocator initialPattern1 10 16) initialPattern 5 9 mapped n
	fpoke (coerce p :: FilePtr Word32) 0xDEAD
	return p

-- | Free a block.
free p = unless (pointer p == nullPtr) $ maskSigint $ withFileLock (path (fileSrc p) ++ ".lock") Exclusive $ const $ void $ do
	fpoke (coerce p :: FilePtr Word32) 0xFEEE
	_free (_free dumbFreer initialPattern1 (truncateOffset 16)) initialPattern ((+ 24) . truncateOffset 10 . subtract 24) p

---------------------------------------------
-- Pointer operations

fpeek p = peek (pointer p)

fpoke p x = poke (pointer p) x

peekPtr :: FilePtr (FilePtr t) -> IO (FilePtr t)
peekPtr p = fpeek p >>= make (fileSrc p)

pokePtr :: FilePtr (FilePtr t) -> FilePtr t -> IO ()
pokePtr p = fpoke p . offset

-- | Use caution with coercions as they can lead to file corruption and memory faults.
coerce :: FilePtr t -> FilePtr u
coerce (FilePtr p o m) = FilePtr (castPtr p) o m

-- | Pointer arithmetic deserves similar caution.
plus :: FilePtr t -> Int -> FilePtr t
plus (FilePtr p off map) n = FilePtr (plusPtr p n) (off + fromIntegral n) map

minus :: FilePtr a -> FilePtr a -> Int
minus (FilePtr _ off (Mapped _ blk)) (FilePtr _ off2 (Mapped _ blk2)) = if blk == blk2 then
		fromIntegral $ off2 - off
	else
		error "File.Mapped.minus: ptrs from different files"

-- | 'toRepr' is a safe version of 'coerce', where the target of the
-- coercion is really the bit representation of the source.
toRepr :: FilePtr t -> FilePtr (Represent t)
toRepr (FilePtr p o m) = FilePtr (castPtr p) o m

---------------------------------------------
-- Signal masking

maskSigint :: IO a -> IO a
#ifndef __WIN32__
maskSigint m = do
	-- Mask those signals
	mask <- getSignalMask
	setSignalMask $ foldr addSignal mask terminatingSignals

	-- Run the thing
	x <- finally
		m
		(do
		-- Attempt to put the signal mask back the way it was
		mask2 <- getSignalMask
		setSignalMask $ foldr (\s -> if inSignalSet s mask then id else deleteSignal s) mask2 terminatingSignals)

	return x

terminatingSignals = [sigINT, sigTERM, sigQUIT]
#else
maskSigint = id
#endif