-- | .2bit format (from the UCSC Genome Browser FAQ)
--
-- A .2bit file stores multiple DNA sequences (up to 4 Gb total) in a
-- compact randomly-accessible format. The file contains masking
-- information as well as the DNA itself.
--
-- The file begins with a 16-byte header containing the following fields:
--
-- - signature - the number 0x1A412743 in the architecture of the machine that created the file
-- - version - zero for now. Readers should abort if they see a version number higher than 0
-- - sequenceCount - the number of sequences in the file
-- - reserved - always zero for now
--
-- All fields are 32 bits unless noted. If the signature value is not as
-- given, the reader program should byte-swap the signature and check if
-- the swapped version matches. If so, all multiple-byte entities in the
-- file will have to be byte-swapped. This enables these binary files to
-- be used unchanged on different architectures.
--
-- The header is followed by a file index, which contains one entry for
-- each sequence. Each index entry contains three fields:
--
-- - nameSize - a byte containing the length of the name field
-- - name - the sequence name itself (in ASCII-compatible byte string), of variable length depending on nameSize
-- - offset - the 32-bit offset of the sequence data relative to the start of the file, not aligned to any 4-byte padding boundary
--
-- The index is followed by the sequence records, which contain nine fields:
--
-- - dnaSize - number of bases of DNA in the sequence
-- - nBlockCount - the number of blocks of Ns in the file (representing unknown sequence)
-- - nBlockStarts - an array of length nBlockCount of 32 bit integers indicating the (0-based) starting position of a block of Ns
-- - nBlockSizes - an array of length nBlockCount of 32 bit integers indicating the length of a block of Ns
-- - maskBlockCount - the number of masked (lower-case) blocks
-- - maskBlockStarts - an array of length maskBlockCount of 32 bit integers indicating the (0-based) starting position of a masked block
-- - maskBlockSizes - an array of length maskBlockCount of 32 bit integers indicating the length of a masked block
-- - reserved - always zero for now
-- - packedDna - the DNA packed to two bits per base, represented as so:
-- T - 00, C - 01, A - 10, G - 11. The first base is in the most
-- significant 2-bit byte; the last base is in the least significant
-- 2 bits. For example, the sequence TCAG is represented as 00011011.
--
-- In this format, it is neither possible nor necessary to store Ns in
-- the main sequence, and one wouldn't expect them to take up space
-- there. However, they do; hard masked sequence is typically stored as
-- many Ts. The sensible way to treat these is probably to just say
-- there are two kinds of implied annotation (repeats and large gaps for
-- a typical genome), which can be interpreted in whatever way fits.
module Bio.TwoBit (
TwoBitFile(..),
openTwoBit,
TwoBitChromosome(..),
tbf_chrnames,
findChrom,
TwoBitSequence'(..),
TwoBitSequence,
Unidrectional,
Bidirectional,
unpackRSRaw,
unpackRS,
unpackRSMasked,
Masking(..),
isSoftMasked,
isHardMasked,
noneMasked,
softMasked,
hardMasked,
bothMasked
) where
import Control.Applicative
import Control.Exception ( Exception(..), throw )
import Control.Monad ( guard )
import Control.Monad.Primitive ( unsafeInlineIO )
import Data.Bits
import Data.Char ( toLower )
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Internal as B ( fromForeignPtr )
import Data.Foldable
import qualified Data.HashMap.Strict as M
import Data.List ( mapAccumL )
import Data.Primitive.Array ( Array, arrayFromList )
import Data.Primitive.PrimArray ( indexPrimArray )
import Data.Word ( byteSwap32, Word8, Word32 )
import Foreign.ForeignPtr ( ForeignPtr, withForeignPtr )
import Foreign.Ptr ( castPtr, plusPtr, Ptr )
import Foreign.Storable ( Storable(..) )
import GHC.Base ( build )
import System.IO.MMap ( mmapFileForeignPtr, Mode(..) )
data TwoBitFile = TBF { tbf_raw :: {-# UNPACK #-} !(ForeignPtr Word8)
, tbf_size :: {-# UNPACK #-} !Int
, tbf_path :: {-# UNPACK #-} !B.ByteString
, tbf_chroms :: {-# UNPACK #-} !(Array TwoBitChromosome)
, tbf_chrmap :: !(M.HashMap B.ByteString TwoBitChromosome) }
tbf_chrnames :: TwoBitFile -> [B.ByteString]
tbf_chrnames = toList . fmap tbc_name . tbf_chroms
-- | Finds a named scaffold in the reference. If it doesn't find the
-- exact name, it will try to compensate for the crazy naming
-- differences between NCBI and UCSC. This doesn't work in general, but
-- is good enough in the common case. In particular, "1" maps to "chr1"
-- and back, "GL000192.1" to "chr1_gl000192_random" and back, and "chrM"
-- to "MT" and back.
findChrom :: B.ByteString -> TwoBitFile -> Maybe TwoBitChromosome
findChrom c TBF{ tbf_chrmap = cs } =
M.lookup c cs
<|> M.lookup ("chr" <> c) cs
<|> ( guard ("chr" `B.isPrefixOf` c) >> M.lookup (B.drop 3 c) cs )
<|> ( guard ("chrM" == c) >> M.lookup "MT" cs )
<|> ( guard ("MT" == c) >> M.lookup "chrM" cs )
<|> ( case filter (\d -> match c (tbc_name d) || match (tbc_name d) c) $ M.elems cs of
[x] -> Just x ; _ -> Nothing )
where
match x y = B.isInfixOf (B.map toLower (B.takeWhile (/= '.') x)) y
data TwoBitChromosome = TBC { tbc_raw :: {-# UNPACK #-} !(ForeignPtr Word8)
, tbc_name :: {-# UNPACK #-} !B.ByteString
, tbc_index :: {-# UNPACK #-} !Int
, tbc_dna_offset :: {-# UNPACK #-} !Word32
, tbc_dna_size :: {-# UNPACK #-} !Word32
-- | Lazily generated sequence in forward direction; the argument is the offset of the first base.
, tbc_fwd_seq :: Int -> TwoBitSequence' Unidrectional
-- | Lazily generated sequence in reverse direction; the argument is the offset of the first base to the
-- right of the beginning. (The first base generated is the complement of the base found at (offset-1).
, tbc_rev_seq :: Int -> TwoBitSequence' Bidirectional }
data TwoBitError = WrongSignature FilePath
| UnsortedBlocks FilePath
| OutOfBounds FilePath Word32 Int
| OverlongSequence FilePath Word32 Word32 Int
deriving Show
instance Exception TwoBitError where
displayException (WrongSignature fp) = "The file " ++ show fp ++ "does not have a .2bit signature."
displayException (UnsortedBlocks fp) = "The N and mask blocks in file " ++ show fp ++ " are not sorted."
displayException (OutOfBounds fp o s) = "Attempted to access offset " ++ show o ++ " in file " ++ show fp ++ " of size " ++ show s ++ "."
displayException (OverlongSequence fp o l s) = "A sequence of length " ++ show l ++ " starting at " ++ show o ++ " in file "
++ show fp ++ " hangs over its end at " ++ show s ++ "."
-- | Brings a 2bit file into memory. The file is mmap'ed, so it will
-- not work on streams that are not actual files. It's also unsafe if
-- the file is concurrently modified in any way.
openTwoBit :: FilePath -> IO TwoBitFile
openTwoBit fp = do (p,o,l) <- mmapFileForeignPtr fp ReadOnly Nothing
if o == 0 then pure $ parseTwoBit fp p l
else fail $ "unexpected: mmapFileForeignPtr returned an offset"
-- | Parses a 2bit file. The @FilePath@ argument is only used in error
-- messages, what is really parsed is the memory block, typically from
-- mmapping the file.
--
-- The workhorse in here is the construction of the 'tbc_fwd_seq' and
-- 'tbc_rev_seq' functions. When called, they first run a binary search
-- on the mask lists, then produce a list of blocks with uniform
-- masking. Both parts of the algorithm are fast and directly use the
-- on-disk data structures.
--
-- In theory, there could be 2bit files in big endian format out there.
-- We nominally support them, but since I've never seen one in the wild,
-- this may well fail in a spectacular way.
parseTwoBit :: FilePath -> ForeignPtr Word8 -> Int -> TwoBitFile
parseTwoBit fp0 raw size
| getW32_ peekUnalnWord32 0 == 0x1A412743 && getW32_ peekUnalnWord32 4 == 0 = kont $ parseEachSeq (getW32_ peekUnalnWord32)
| getW32_ peekUnalnWord32Swap 0 == 0x1A412743 && getW32_ peekUnalnWord32Swap 4 == 0 = kont $ parseEachSeq (getW32_ peekUnalnWord32Swap)
| otherwise = throw $ WrongSignature fp0
where
kont sqs = TBF raw size (B.pack fp0) (arrayFromList sqs) (M.fromList $ map (liftA2 (,) tbc_name id) sqs)
getW32_ f o | o + 4 >= fromIntegral size = throw $ OutOfBounds fp0 o size
| otherwise = unsafeInlineIO $ withForeignPtr raw $ \p -> f (plusPtr p (fromIntegral o))
parseEachSeq :: (Word32 -> Word32) -> [TwoBitChromosome]
parseEachSeq getW32 = snd $ mapAccumL (parseOneSeq getW32) 16 [0 .. fromIntegral (getW32 8) -1]
parseOneSeq getW32 off nseq =
if packedDnaOff + shiftR (dnasize+3) 2 > fromIntegral size
then throw $ OverlongSequence fp0 packedDnaOff dnasize size
else (off + 5 + nmsize, TBC raw name nseq packedDnaOff dnasize unfoldSeqFwd unfoldSeqRev)
where
!nmsize = unsafeInlineIO $ withForeignPtr raw $ \p -> fromIntegral <$> peekElemOff p off
!name = B.fromForeignPtr raw (off+1) nmsize
!offset = getW32 . fromIntegral $ off + 1 + nmsize
!dnasize = getW32 $ offset
!nBlockCount = getW32 $ offset + 4
!mBlockCount = getW32 $ offset + 8 + 8*nBlockCount
!packedDnaOff = offset + 16 + 8 * (nBlockCount+mBlockCount)
-- Valid blocks are numbered 1..max; there are virtual guard blocks at indices 0 and (max+1), which make the later
-- algorithms much cleaner
n_block, m_block :: Word32 -> Block
n_block i | i == 0 = B 0 0 i
| i > nBlockCount = B maxBound maxBound i
| otherwise = B a (a+b) i
where
!a = getW32 $ offset+4 + 4*i
!b = getW32 $ offset+4 + 4*(i+nBlockCount)
m_block i | i == 0 = B 0 0 i
| i > mBlockCount = B maxBound maxBound i
| otherwise = B a (a+b) i
where
!a = getW32 $ offset+8 + 8*nBlockCount + 4*i
!b = getW32 $ offset+8 + 8*nBlockCount + 4*(i+mBlockCount)
unfoldSeqFwd :: Int -> TwoBitSequence' Unidrectional
unfoldSeqFwd chroff = unfoldSeqFwd' (search n_block nBlockCount) (search m_block mBlockCount)
(fromIntegral chroff) (packedDnaOff * 4 + fromIntegral chroff)
where
trim b = b { start_offset = max (fromIntegral chroff) (start_offset b) }
-- finds the smallest index such that the block end(!) is larger than 'chroff'
search f num = trim . f $ go 0 (num+1)
where
go a b | a == b = a
| end_offset (f m) > fromIntegral chroff = go a m
| otherwise = go (m+1) b
where
m = div (a + b) 2
unfoldSeqFwd' :: Block -> Block -> Word32 -> Word32 -> TwoBitSequence' Unidrectional
unfoldSeqFwd' nb@(B nstart nend _) mb@(B mstart mend _) !chroff !fileoff
| chroff >= dnasize = RefEnd
| chroff > mstart || chroff > nstart = throw (UnsortedBlocks fp0)
| chroff < nstart && chroff < mstart = advance noneMasked $ min dnasize $ min nstart mstart
| chroff < mstart = advance hardMasked $ min dnasize $ min nend mstart
| chroff < nstart = advance softMasked $ min dnasize $ min mend nstart
| otherwise = advance bothMasked $ min dnasize $ min nend mend
where
advance m x = SomeSeq m raw (fromIntegral fileoff) (fromIntegral $ x - chroff) $
unfoldSeqFwd' (advanceB x n_block nb) (advanceB x m_block mb) x (fileoff+x-chroff)
advanceB :: Word32 -> (Word32 -> Block) -> Block -> Block
advanceB x f (B start end i)
| x <= start = B start end i
| x < end = B x end i
| otherwise = f (i+1)
unfoldSeqRev :: Int -> TwoBitSequence' Bidirectional
unfoldSeqRev chroff = unfoldSeqRev' (search n_block nBlockCount) (search m_block mBlockCount)
(fromIntegral chroff) (packedDnaOff * 4 + fromIntegral chroff)
where
trim b = b { end_offset = min (fromIntegral chroff) (end_offset b) }
-- finds the largest index such that the block start is smaller than chroff
search f num = trim . f $ go 0 (num+1)
where
go a b | a == b = a
| start_offset (f m) < fromIntegral chroff = go m b
| otherwise = go a (m-1)
where
m = div (a + b + 1) 2
unfoldSeqRev' :: Block -> Block -> Word32 -> Word32 -> TwoBitSequence' Bidirectional
unfoldSeqRev' nb@(B nstart nend _) mb@(B mstart mend _) !chroff !fileoff
| chroff <= 0 = RefEnd
| chroff < mend || chroff < nend = throw (UnsortedBlocks fp0)
| chroff > nend && chroff > mend = advance noneMasked $ max nend mend
| chroff > mend = advance hardMasked $ max nstart mend
| chroff > nend = advance softMasked $ max mstart nend
| otherwise = advance bothMasked $ max nstart mstart
where
advance m x = SomeSeq m raw (fromIntegral fileoff) (fromIntegral x - fromIntegral chroff) $
unfoldSeqRev' (advanceB x n_block nb) (advanceB x m_block mb) x (fileoff+x-chroff)
advanceB :: Word32 -> (Word32 -> Block) -> Block -> Block
advanceB x f (B start end i)
| x >= end = B start end i
| x > start = B start x i
| otherwise = f (i-1)
data Block = B { start_offset :: !Word32
, end_offset :: !Word32
, block_number :: !Word32 }
deriving (Show, Eq, Ord)
-- | 2bit supports two kinds of masking, typically rendered as lowercase
-- letters ('MaskSoft') and Ns ('MaskHard'). They can overlap
-- ('MaskBoth'), and even the hard masking has underlying sequence
-- (which is normally ignored).
newtype Masking = Masking Word8 deriving (Eq, Ord)
instance Show Masking where
show (Masking 0) = "None"
show (Masking 1) = "Soft"
show (Masking 2) = "Hard"
show (Masking _) = "Both"
instance Read Masking where
readsPrec _ s = [ (Masking m,s') | (w,s') <- lex s
, m <- case w of "None" -> [0]
"Soft" -> [1]
"Hard" -> [2]
"Both" -> [3]
_ -> [ ] ]
instance Semigroup Masking where
Masking a <> Masking b = Masking (a .|. b)
instance Monoid Masking where
mempty = Masking 0
mappend = (<>)
instance Enum Masking where
toEnum = Masking . toEnum
fromEnum (Masking m) = fromEnum m
instance Bounded Masking where
minBound = Masking 0
maxBound = Masking 3
isSoftMasked, isHardMasked :: Masking -> Bool
isSoftMasked (Masking m) = testBit m 0
isHardMasked (Masking m) = testBit m 1
noneMasked, softMasked, hardMasked, bothMasked :: Masking
noneMasked = Masking 0
softMasked = Masking 1
hardMasked = Masking 2
bothMasked = Masking 3
-- | This is a (piece of a) reference sequence. It consists of
-- stretches with uniform masking.
--
-- The offset is stored as a 'Word'. This is done because on a 32 bit
-- platform, every bit counts. This limits the genome to approximately
-- four gigabases, which would be a file of about one gigabyte. That's
-- just about enough to work with the human genome. On a 64 bit
-- platform, the file format itself imposes a limit of four gigabytes,
-- or about 16 gigabases in total.
--
-- If length is zero, the piece is empty and the mask, pointer, and
-- offset fields may not be valid. If length is positive, ptr+offset
-- points at the first base of the piece. If length is negative,
-- ptr+offset points just past the end of the piece, ptr+offset+length
-- points to the first base of the piece, and the sequence in meant to
-- be reverse complemented.
--
-- In a 'TwoBitSequence', length must not be negative. In a
-- @TwoBitSequence' Bidirectional@, length can be positive or negative.
data TwoBitSequence' dir = SomeSeq {-# UNPACK #-} !Masking -- ^ how is it masked?
{-# UNPACK #-} !(ForeignPtr Word8) -- ^ primitive bases in 2bit encoding: [0..3] = TCAG
{-# UNPACK #-} !Word -- ^ offset in bases(!)
{-# UNPACK #-} !Int -- ^ length in bases
(TwoBitSequence' dir)
| RefEnd
data Unidrectional
data Bidirectional
type TwoBitSequence = TwoBitSequence' Unidrectional
instance Show (TwoBitSequence' dir) where
showsPrec _ (SomeSeq m _ _ l r) = (++) "SomeSeq " . shows m . (:) ' ' . shows l . (++) " $ " . shows r
showsPrec _ RefEnd = (++) "RefEnd"
-- | Unpacks a reference sequence into a (very long) list of bytes.
-- Each byte contains the nucleotide in bits 0 and 1 with valjues 0..3
-- corresponding to "TCAG", and the soft and hard mask bits in bits 2
-- and 3, respectively.
unpackRSRaw :: TwoBitSequence' dir -> [Word8]
unpackRSRaw rs = build (\c n -> unpackRSFB c n rs)
{-# INLINE unpackRSRaw #-}
unpackRSFB :: (Word8 -> b -> b) -> b -> TwoBitSequence' dir -> b
unpackRSFB cons nil = go0
where
go0 RefEnd = nil
go0 (SomeSeq (Masking msk) raw off0 len0 rs) | len0 >= 0 = go off0 len0
where
go !off !len = if len == 0 then go0 rs else code `cons` go (off+1) (len-1)
where
!byteoff = fromIntegral $ off `shiftR` 2
!bitoff = fromIntegral $ off .&. 3
!byte = unsafeInlineIO $ withForeignPtr raw (`peekByteOff` byteoff)
!code = shiftR byte (6 - 2 * bitoff ) .&. 3 .|. shiftL msk 2
go0 (SomeSeq (Masking msk) raw off0 len0 rs) = go off0 (-len0)
where
go !off !len = if len == 0 then go0 rs else xor 2 code `cons` go (off-1) (len-1)
where
!byteoff = fromIntegral $ (off-1) `shiftR` 2
!bitoff = fromIntegral $ (off-1) .&. 3
!byte = unsafeInlineIO $ withForeignPtr raw (`peekByteOff` byteoff)
!code = shiftR byte (6 - 2 * bitoff ) .&. 3 .|. shiftL msk 2
{-# INLINE [0] unpackRSFB #-}
-- | Unpacks a reference sequence into a (very long) list of ASCII
-- characters. Hard masked nucleotides become the letter 'N', others
-- become "TCAG".
unpackRS :: TwoBitSequence' dir -> [Word8]
unpackRS = map (indexPrimArray chars . fromIntegral) . unpackRSRaw
where
!chars = [84,67,65,71,84,67,65,71,78,78,78,78,78,78,78,78] -- "TCAGTCAGNNNNNNNN"
{-# INLINE unpackRS #-}
-- | Unpacks a reference sequence into a list of ASCII characters,
-- interpreting masking in the customary way. Specifically, hard
-- masking produces Ns, soft masking produces lower case letters, and
-- dual masking produces lower case Ns.
unpackRSMasked :: TwoBitSequence' dir -> [Word8]
unpackRSMasked = map (indexPrimArray chars . fromIntegral) . unpackRSRaw
where
!chars = [84,67,65,71,116,99,97,103,78,78,78,78,110,110,110,110] -- "TCAGtcagNNNNnnnn"
{-# INLINE unpackRSMasked #-}
-- | Reads a 32 bit word from an address, which doesn't need to be
-- aligned. The byte order used is unspecified.
peekUnalnWord32 :: Ptr a -> IO Word32
-- | Equivalent to peekUnalnWord32 followed by a byte swap.
peekUnalnWord32Swap :: Ptr a -> IO Word32
-- List of known architectures that efficiently support unaligned accesses.
#if defined(i386_HOST_ARCH) || defined(x86_64_HOST_ARCH) \
|| defined(powerpc64le_HOST_ARCH) || ((defined(arm_HOST_ARCH) \
|| defined(aarch64_HOST_ARCH)) && defined(__ARM_FEATURE_UNALIGNED)) \
|| defined(powerpc_HOST_ARCH) || defined(powerpc64_HOST_ARCH)
peekUnalnWord32 = peek . castPtr
peekUnalnWord32Swap = fmap byteSwap32 . peek . castPtr
#else
peekUnalnWord32 p = do
x <- fromIntegral <$> peekWord8 (plusPtr p 0)
y <- fromIntegral <$> peekWord8 (plusPtr p 1)
z <- fromIntegral <$> peekWord8 (plusPtr p 2)
w <- fromIntegral <$> peekWord8 (plusPtr p 3)
return $! x .|. unsafeShiftL y 8 .|. unsafeShiftL z 16 .|. unsafeShiftL w 24
peekUnalnWord32Swap p = do
x <- fromIntegral <$> peekWord8 (plusPtr p 0)
y <- fromIntegral <$> peekWord8 (plusPtr p 1)
z <- fromIntegral <$> peekWord8 (plusPtr p 2)
w <- fromIntegral <$> peekWord8 (plusPtr p 3)
return $! w .|. unsafeShiftL z 8 .|. unsafeShiftL y 16 .|. unsafeShiftL x 24
peekWord8 :: Ptr a -> IO Word8
peekWord8 = peek . castPtr
#endif