content: module Compress(main) where
  import Prelude
  import Data.Map as M
  import Data.Char
  import System.Environment
  import System.IO
  --import Debug.Trace
  
  --
  -- A simple LZW compressor
  -- It uses a dictionary of maximum size 4096.
  -- The 12 bit code words are encoded with 2 code words in 3 bytes.
  -- It can only compress input of printable ASCII + '\n'
  
  type Table = M.Map [Char] Int
  
  -- Don't change this lightly.
  -- It needs to be coordinated with the decompressor transducer in eval.c
  maxDict :: Int
  maxDict = 4096
  
  toChar :: Int -> Char
  toChar i = chr (i + 32)
  
  (!) :: Table -> [Char] -> Int
  (!) t s =
    case M.lookupBy compare s t of
      Nothing -> undefined -- error $ "(!): " ++ showString s
      Just i -> i
  
  compress :: Table -> [Char] -> [Char] -> [Int]
  compress t [] p = [ t ! p ]
  compress t (c:cs) p =
    let p' = p ++ [c]
        s = M.size t
        t' = if s < maxDict then M.insertBy compare p' s t else t
    in
  --      trace ("compress " ++ showString p') $
  --      trace show (M.toList t)) $
        case M.lookupBy compare p' t of
          Just _ ->
  --          trace "found" $
            compress t cs p'
          Nothing ->
  --          trace ("not found p=" ++ show p ++ " " ++ show (M.lookupBy compare p t)) $
            (t ! p) : compress t' cs [c]
  
  -- Initial table is ' ' .. '~', and '\n'
  initTable :: Table
  initTable = M.fromListBy compare $ [([toChar c], c) | c <- [0..94] ] ++ [("\n", 95)]
  
  toBytes :: [Int] -> [Int]
  toBytes [] = []
  toBytes [i] = [i `rem` 256, i `quot` 256, 0]
  toBytes (i1:i2:is) =
    let i = i1 + 4096*i2
        b1 = i `rem` 256
        b2 = (i `quot` 256) `rem` 256
        b3 = i `quot` (256*256)
    in  b1 : b2 : b3 : toBytes is
  
  bad :: Char -> Bool
  bad c = not (isPrint c || c == '\n')
  
  main :: IO ()
  main = do
    args <- getArgs
    let (ifn, ofn) = case args of { [a,b] -> (a,b); _ -> error "Usage: Compress in-file out-file" }
    ifile <- openFile ifn ReadMode
    f <- hGetContents ifile
    when (any bad f) $
      error "Non-printable ASCII in input"
    let bs = compress initTable f []
    ofile <- openBinaryFile ofn WriteMode
    hPutStr ofile $ 'Z' : map chr (toBytes bs)
    hClose ifile
    hClose ofile
  