Сохранение кода в картинке

lab3/src/Lib.hs

@@ -1,18 +1,19 @@
 module Lib
-    ( 
-        createAlphabetFromText,
-        encryptCaesar,
-        decryptCaesar,
-        textToBits,
-        bitsToText,
-        encodePixel
-    ) where
+    -- ( 
+        -- createAlphabetFromText,
+        -- encryptCaesar,
+        -- decryptCaesar,
+        -- textToBits,
+        -- bitsToText,
+        -- encodePixel
+    -- ) 
+    where
 
 import Codec.Picture
 
 import Data.Word (Word8)
 import Data.Char (ord, chr)
-import Data.Bits (testBit)
+import Data.Bits (testBit, shiftL, complement, (.|.), (.&.))
 import qualified Data.Vector.Unboxed as VU
 
 createAlphabetFromText :: String -> [Char]
@@ -46,14 +47,34 @@ textToBits text = VU.fromList $ concatMap charToBits text
 charToBits :: Char -> [Int]
 charToBits c = [if testBit (ord c) i then 1 else 0 | i <- [7,6..0]]
 
+intToBits :: Int -> [Int]
+intToBits n = [if testBit n i then 1 else 0 | i <- [31,30..0]]
+
 bitsToText :: VU.Vector Int -> String
 bitsToText bits
     | VU.null bits = []
     | otherwise = (chr $ bitsToInt (VU.take 8 bits)) : bitsToText (VU.drop 8 bits)
 
 bitsToInt :: VU.Vector Int -> Int
-bitsToInt charBits = 
-    sum [bit * (2 ^ index) | (bit, index) <- zip (VU.toList charBits) [7 :: Int,6..0]]
+bitsToInt bits = 
+    sum [bit * (2 ^ index) | (bit, index) <- zip (VU.toList bits) [len,(len - 1)..0]]
+    where
+        len = VU.length bits - 1
+
+setLastBits :: VU.Vector Int -> VU.Vector Int -> VU.Vector Int
+setLastBits byte newBits = byte VU.// updates
+  where
+    newBitsLastIndex = VU.length newBits - 1
+    updates = [(7 - i, newBits VU.! (newBitsLastIndex - i)) | i <- [0..newBitsLastIndex]]
+
+intToWord8 :: Int -> Word8
+intToWord8 x = fromIntegral x
+
+word8ToInt :: Word8 -> Int
+word8ToInt x = fromIntegral x
+
+createMask :: Int -> Int
+createMask shift = shiftL (complement 0) shift .&. complement 0
 
 encodePixel :: Int -> Image PixelRGB8 -> VU.Vector Int -> Int -> Int -> PixelRGB8
 encodePixel bitsPerByte img bits x y = PixelRGB8 newR newG newB
@@ -64,14 +85,13 @@ encodePixel bitsPerByte img bits x y = PixelRGB8 newR newG newB
         startPos = index * 3 * bitsPerByte
         pixelBits = VU.slice startPos (3 * bitsPerByte) bits
 
-        intToWord8 :: Int -> Word8
-        intToWord8 x = fromIntegral x
+        bitsIntR = bitsToInt $ VU.slice 0 bitsPerByte pixelBits
+        bitsIntG = bitsToInt $ VU.slice bitsPerByte bitsPerByte pixelBits
+        bitsIntB = bitsToInt $ VU.slice (2 * bitsPerByte) bitsPerByte pixelBits
 
-        bitsR = intToWord8 $ (VU.take bitsPerByte pixelBits) VU.! 0 
-        bitsG = bitsToInt $ VU.take bitsPerByte $ VU.drop bitsPerByte pixelBits
-        bitsB = bitsToInt $ VU.drop (2 * bitsPerByte) pixelBits
+        mask = createMask bitsPerByte
 
         PixelRGB8 r g b = pixelAt img x y
-        newR = bitsR
-        newG = intToWord8 bitsG
-        newB = b
+        newR = intToWord8 $ ((word8ToInt r) .&. mask) .|. bitsIntR
+        newG = intToWord8 $ ((word8ToInt g) .&. mask) .|. bitsIntG
+        newB = intToWord8 $ ((word8ToInt b) .&. mask) .|. bitsIntB