/*
 * Copyright (C) 2014 Michael Brown <mbrown@fensystems.co.uk>.
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; either version 2 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301, USA.
 */

#include "mscompress.h"

#include <grub/misc.h>

 /** Number of aligned offset codes */
#define LZX_ALIGNOFFSET_CODES 8

/** Aligned offset code length (in bits) */
#define LZX_ALIGNOFFSET_BITS 3

/** Number of pretree codes */
#define LZX_PRETREE_CODES 20

/** Pretree code length (in bits) */
#define LZX_PRETREE_BITS 4

/** Number of literal main codes */
#define LZX_MAIN_LIT_CODES 256

/** Number of position slots */
#define LZX_POSITION_SLOTS 30

/** Number of main codes */
#define LZX_MAIN_CODES (LZX_MAIN_LIT_CODES + (8 * LZX_POSITION_SLOTS))

/** Number of length codes */
#define LZX_LENGTH_CODES 249

/** Block type length (in bits) */
#define LZX_BLOCK_TYPE_BITS 3

/** Default block length */
#define LZX_DEFAULT_BLOCK_LEN 32768

/** Number of repeated offsets */
#define LZX_REPEATED_OFFSETS 3

/** Don't ask */
#define LZX_WIM_MAGIC_FILESIZE 12000000

/** Block types */
enum lzx_block_type
{
	/** Verbatim block */
	LZX_BLOCK_VERBATIM = 1,
	/** Aligned offset block */
	LZX_BLOCK_ALIGNOFFSET = 2,
	/** Uncompressed block */
	LZX_BLOCK_UNCOMPRESSED = 3,
};

/** An LZX input stream */
struct lzx_input_stream
{
	/** Data */
	const grub_uint8_t* data;
	/** Length */
	grub_size_t len;
	/** Offset within stream */
	grub_size_t offset;
};

/** An LZX output stream */
struct lzx_output_stream
{
	/** Data, or NULL */
	grub_uint8_t* data;
	/** Offset within stream */
	grub_size_t offset;
	/** End of current block within stream */
	grub_size_t threshold;
};

/** LZX decompressor */
struct lzx
{
	/** Input stream */
	struct lzx_input_stream input;
	/** Output stream */
	struct lzx_output_stream output;
	/** Accumulator */
	grub_uint32_t accumulator;
	/** Number of bits in accumulator */
	unsigned int bits;
	/** Block type */
	enum lzx_block_type block_type;
	/** Repeated offsets */
	unsigned int repeated_offset[LZX_REPEATED_OFFSETS];

	/** Aligned offset Huffman alphabet */
	struct huffman_alphabet alignoffset;
	/** Aligned offset raw symbols
	 *
	 * Must immediately follow the aligned offset Huffman
	 * alphabet.
	 */
	huffman_raw_symbol_t alignoffset_raw[LZX_ALIGNOFFSET_CODES];
	/** Aligned offset code lengths */
	grub_uint8_t alignoffset_lengths[LZX_ALIGNOFFSET_CODES];

	/** Pretree Huffman alphabet */
	struct huffman_alphabet pretree;
	/** Pretree raw symbols
	 *
	 * Must immediately follow the pretree Huffman alphabet.
	 */
	huffman_raw_symbol_t pretree_raw[LZX_PRETREE_CODES];
	/** Preetree code lengths */
	grub_uint8_t pretree_lengths[LZX_PRETREE_CODES];

	/** Main Huffman alphabet */
	struct huffman_alphabet main;
	/** Main raw symbols
	 *
	 * Must immediately follow the main Huffman alphabet.
	 */
	huffman_raw_symbol_t main_raw[LZX_MAIN_CODES];
	/** Main code lengths */
	GRUB_PACKED_START
	struct
	{
		/** Literals */
		grub_uint8_t literals[LZX_MAIN_LIT_CODES];
		/** Remaining symbols */
		grub_uint8_t remainder[LZX_MAIN_CODES - LZX_MAIN_LIT_CODES];
	} main_lengths;
	GRUB_PACKED_END
	/** Length Huffman alphabet */
	struct huffman_alphabet length;
	/** Length raw symbols
	 *
	 * Must immediately follow the length Huffman alphabet.
	 */
	huffman_raw_symbol_t length_raw[LZX_LENGTH_CODES];
	/** Length code lengths */
	grub_uint8_t length_lengths[LZX_LENGTH_CODES];
};

 /** Base positions, indexed by position slot */
static unsigned int lzx_position_base[LZX_POSITION_SLOTS];

/**
 * Calculate number of footer bits for a given position slot
 *
 * @v position_slot  Position slot
 * @ret footer_bits   Number of footer bits
 */
static inline unsigned int lzx_footer_bits(unsigned int position_slot)
{
	if (position_slot < 2)
		return 0;
	else if (position_slot < 38)
		return ((position_slot / 2) - 1);
	else
		return 17;
}

/**
 * Attempt to accumulate bits from LZX bitstream
 *
 * @v lzx    Decompressor
 * @v bits    Number of bits to accumulate
 * @v norm_value  Accumulated value (normalised to 16 bits)
 *
 * Note that there may not be sufficient accumulated bits in the
 * bitstream; callers must check that sufficient bits are available
 * before using the value.
 */
static int
lzx_accumulate(struct lzx* lzx, unsigned int bits)
{
	const grub_uint16_t* src16;

	/* Accumulate more bits if required */
	if ((lzx->bits < bits) && (lzx->input.offset < lzx->input.len))
	{
		src16 = (void*)((char*)lzx->input.data + lzx->input.offset);
		lzx->input.offset += sizeof(*src16);
		lzx->accumulator |= (*src16 << (16 - lzx->bits));
		lzx->bits += 16;
	}

	return (lzx->accumulator >> 16);
}

/**
 * Consume accumulated bits from LZX bitstream
 *
 * @v lzx    Decompressor
 * @v bits    Number of bits to consume
 * @ret rc    Return status code
 */
static int lzx_consume(struct lzx* lzx, unsigned int bits)
{
	/* Fail if insufficient bits are available */
	if (lzx->bits < bits)
	{
		return -1;
	}

	/* Consume bits */
	lzx->accumulator <<= bits;
	lzx->bits -= bits;

	return 0;
}

/**
 * Get bits from LZX bitstream
 *
 * @v lzx    Decompressor
 * @v bits    Number of bits to fetch
 * @ret value    Value, or negative error
 */
static int lzx_getbits(struct lzx* lzx, unsigned int bits)
{
	int norm_value;
	int rc;

	/* Accumulate more bits if required */
	norm_value = lzx_accumulate(lzx, bits);

	/* Consume bits */
	if ((rc = lzx_consume(lzx, bits)) != 0)
		return rc;

	return (norm_value >> (16 - bits));
}

/**
 * Align LZX bitstream for byte access
 *
 * @v lzx    Decompressor
 * @v bits    Minimum number of padding bits
 * @ret rc    Return status code
 */
static int lzx_align(struct lzx* lzx, unsigned int bits)
{
	int pad;

	/* Get padding bits */
	pad = lzx_getbits(lzx, bits);
	if (pad < 0)
		return pad;

	/* Consume all accumulated bits */
	lzx_consume(lzx, lzx->bits);

	return 0;
}

/**
 * Get bytes from LZX bitstream
 *
 * @v lzx    Decompressor
 * @v data    Data buffer, or NULL
 * @v len    Length of data buffer
 * @ret rc    Return status code
 */
static int lzx_getbytes(struct lzx* lzx, void* data, grub_size_t len)
{
	/* Sanity check */
	if ((lzx->input.offset + len) > lzx->input.len)
	{
		return -1;
	}

	/* Copy data */
	if (data)
		grub_memmove(data, (lzx->input.data + lzx->input.offset), len);
	lzx->input.offset += len;

	return 0;
}

/**
 * Decode LZX Huffman-coded symbol
 *
 * @v lzx    Decompressor
 * @v alphabet    Huffman alphabet
 * @ret raw    Raw symbol, or negative error
 */
static int lzx_decode(struct lzx* lzx, struct huffman_alphabet* alphabet)
{
	struct huffman_symbols* sym;
	int huf;
	int rc;

	/* Accumulate sufficient bits */
	huf = lzx_accumulate(lzx, HUFFMAN_BITS);
	if (huf < 0)
		return huf;

	/* Decode symbol */
	sym = grub_huffman_sym(alphabet, huf);

	/* Consume bits */
	if ((rc = lzx_consume(lzx, grub_huffman_len(sym))) != 0)
		return rc;

	return grub_huffman_raw(sym, huf);
}

/**
 * Generate Huffman alphabet from raw length table
 *
 * @v lzx    Decompressor
 * @v count    Number of symbols
 * @v bits    Length of each length (in bits)
 * @v lengths    Lengths table to fill in
 * @v alphabet    Huffman alphabet to fill in
 * @ret rc    Return status code
 */
static int
lzx_raw_alphabet(struct lzx* lzx, unsigned int count,
	unsigned int bits, grub_uint8_t* lengths,
	struct huffman_alphabet* alphabet)
{
	unsigned int i;
	int len;
	int rc;

	/* Read lengths */
	for (i = 0; i < count; i++)
	{
		len = lzx_getbits(lzx, bits);
		if (len < 0)
			return len;
		lengths[i] = len;
	}

	/* Generate Huffman alphabet */
	if ((rc = grub_huffman_alphabet(alphabet, lengths, count)) != 0)
		return rc;

	return 0;
}

/**
 * Generate pretree
 *
 * @v lzx    Decompressor
 * @v count    Number of symbols
 * @v lengths    Lengths table to fill in
 * @ret rc    Return status code
 */
static int
lzx_pretree(struct lzx* lzx, unsigned int count, grub_uint8_t* lengths)
{
	unsigned int i;
	unsigned int length;
	int dup = 0;
	int code;
	int rc;

	/* Generate pretree alphabet */
	if ((rc = lzx_raw_alphabet(lzx, LZX_PRETREE_CODES,
		LZX_PRETREE_BITS, lzx->pretree_lengths,
		&lzx->pretree)) != 0)
		return rc;

	/* Read lengths */
	for (i = 0; i < count; i++)
	{
		if (dup)
		{
			/* Duplicate previous length */
			lengths[i] = lengths[i - 1];
			dup--;
		}
		else
		{
			/* Get next code */
			code = lzx_decode(lzx, &lzx->pretree);
			if (code < 0)
				return code;

			/* Interpret code */
			if (code <= 16)
			{
				length = ((lengths[i] - code + 17) % 17);
			}
			else if (code == 17)
			{
				length = 0;
				dup = lzx_getbits(lzx, 4);
				if (dup < 0)
					return dup;
				dup += 3;
			}
			else if (code == 18)
			{
				length = 0;
				dup = lzx_getbits(lzx, 5);
				if (dup < 0)
					return dup;
				dup += 19;
			}
			else if (code == 19)
			{
				length = 0;
				dup = lzx_getbits(lzx, 1);
				if (dup < 0)
					return dup;
				dup += 3;
				code = lzx_decode(lzx, &lzx->pretree);
				if (code < 0)
					return code;
				length = ((lengths[i] - code + 17) % 17);
			}
			else
			{
				return -1;
			}
			lengths[i] = length;
		}
	}

	/* Sanity check */
	if (dup)
	{
		return -1;
	}

	return 0;
}

/**
 * Generate aligned offset Huffman alphabet
 *
 * @v lzx    Decompressor
 * @ret rc    Return status code
 */
static int lzx_alignoffset_alphabet(struct lzx* lzx)
{
	int rc;
	/* Generate aligned offset alphabet */
	if ((rc = lzx_raw_alphabet(lzx, LZX_ALIGNOFFSET_CODES, LZX_ALIGNOFFSET_BITS,
		lzx->alignoffset_lengths, &lzx->alignoffset)) != 0)
		return rc;

	return 0;
}

/**
 * Generate main Huffman alphabet
 *
 * @v lzx    Decompressor
 * @ret rc    Return status code
 */
static int lzx_main_alphabet(struct lzx* lzx)
{
	int rc;

	/* Generate literal symbols pretree */
	if ((rc = lzx_pretree(lzx, LZX_MAIN_LIT_CODES, lzx->main_lengths.literals)) != 0)
	{
		return rc;
	}

	/* Generate remaining symbols pretree */
	if ((rc = lzx_pretree(lzx, (LZX_MAIN_CODES - LZX_MAIN_LIT_CODES),
		lzx->main_lengths.remainder)) != 0)
	{
		return rc;
	}

	/* Generate Huffman alphabet */
	if ((rc = grub_huffman_alphabet(&lzx->main, lzx->main_lengths.literals,
		LZX_MAIN_CODES)) != 0)
	{
		return rc;
	}

	return 0;
}

/**
 * Generate length Huffman alphabet
 *
 * @v lzx    Decompressor
 * @ret rc    Return status code
 */
static int lzx_length_alphabet(struct lzx* lzx)
{
	int rc;

	/* Generate pretree */
	if ((rc = lzx_pretree(lzx, LZX_LENGTH_CODES, lzx->length_lengths)) != 0)
	{
		return rc;
	}

	/* Generate Huffman alphabet */
	if ((rc = grub_huffman_alphabet(&lzx->length, lzx->length_lengths,
		LZX_LENGTH_CODES)) != 0)
	{
		return rc;
	}

	return 0;
}

/**
 * Process LZX block header
 *
 * @v lzx    Decompressor
 * @ret rc    Return status code
 */
static int lzx_block_header(struct lzx* lzx)
{
	grub_size_t block_len;
	int block_type;
	int default_len;
	int len_high;
	int len_low;
	int rc;

	/* Get block type */
	block_type = lzx_getbits(lzx, LZX_BLOCK_TYPE_BITS);
	if (block_type < 0)
		return block_type;
	lzx->block_type = block_type;

	/* Check block length */
	default_len = lzx_getbits(lzx, 1);
	if (default_len < 0)
		return default_len;
	if (default_len)
	{
		block_len = LZX_DEFAULT_BLOCK_LEN;
	}
	else
	{
		len_high = lzx_getbits(lzx, 8);
		if (len_high < 0)
			return len_high;
		len_low = lzx_getbits(lzx, 8);
		if (len_low < 0)
			return len_low;
		block_len = ((len_high << 8) | len_low);
	}
	lzx->output.threshold = (lzx->output.offset + block_len);

	/* Handle block type */
	switch (block_type)
	{
	case LZX_BLOCK_ALIGNOFFSET:
		/* Generated aligned offset alphabet */
		if ((rc = lzx_alignoffset_alphabet(lzx)) != 0)
			return rc;
		/* Fall through */
	case LZX_BLOCK_VERBATIM:
		/* Generate main alphabet */
		if ((rc = lzx_main_alphabet(lzx)) != 0)
			return rc;
		/* Generate lengths alphabet */
		if ((rc = lzx_length_alphabet(lzx)) != 0)
			return rc;
		break;
	case LZX_BLOCK_UNCOMPRESSED:
		/* Align input stream */
		if ((rc = lzx_align(lzx, 1)) != 0)
			return rc;
		/* Read new repeated offsets */
		if ((rc = lzx_getbytes(lzx, &lzx->repeated_offset,
			sizeof(lzx->repeated_offset))) != 0)
			return rc;
		break;
	default:
		return -1;
	}

	return 0;
}

/**
 * Process uncompressed data
 *
 * @v lzx    Decompressor
 * @ret rc    Return status code
 */
static int lzx_uncompressed(struct lzx* lzx)
{
	void* data;
	grub_size_t len;
	int rc;

	/* Copy bytes */
	data = lzx->output.data ? (lzx->output.data + lzx->output.offset) : ((void*)0);
	len = (lzx->output.threshold - lzx->output.offset);
	if ((rc = lzx_getbytes(lzx, data, len)) != 0)
		return rc;

	/* Align input stream */
	if (len % 2)
		lzx->input.offset++;

	return 0;
}

/**
 * Process an LZX token
 *
 * @v lzx    Decompressor
 * @ret rc    Return status code
 *
 * Variable names are chosen to match the LZX specification
 * pseudo-code.
 */
static int lzx_token(struct lzx* lzx)
{
	unsigned int length_header;
	unsigned int position_slot;
	unsigned int offset_bits;
	unsigned int i;
	grub_size_t match_offset;
	grub_size_t match_length;
	int verbatim_bits;
	int aligned_bits;
	int lzx_main;
	int length;
	grub_uint8_t* copy;

	/* Get lzx_main symelse*/
	lzx_main = lzx_decode(lzx, &lzx->main);
	if (lzx_main < 0)
		return lzx_main;

	/* Check for literals */
	if (lzx_main < LZX_MAIN_LIT_CODES)
	{
		if (lzx->output.data)
			lzx->output.data[lzx->output.offset] = lzx_main;
		lzx->output.offset++;
		return 0;
	}
	lzx_main -= LZX_MAIN_LIT_CODES;

	/* Calculate the match length */
	length_header = (lzx_main & 7);
	if (length_header == 7)
	{
		length = lzx_decode(lzx, &lzx->length);
		if (length < 0)
			return length;
	}
	else
	{
		length = 0;
	}
	match_length = (length_header + 2 + length);

	/* Calculate the position slot */
	position_slot = (lzx_main >> 3);
	if (position_slot < LZX_REPEATED_OFFSETS)
	{
		/* Repeated offset */
		match_offset = lzx->repeated_offset[position_slot];
		lzx->repeated_offset[position_slot] = lzx->repeated_offset[0];
		lzx->repeated_offset[0] = match_offset;
	}
	else
	{
		/* Non-repeated offset */
		offset_bits = lzx_footer_bits(position_slot);
		if ((lzx->block_type == LZX_BLOCK_ALIGNOFFSET) &&
			(offset_bits >= 3))
		{
			verbatim_bits = lzx_getbits(lzx, (offset_bits - 3));
			if (verbatim_bits < 0)
				return verbatim_bits;
			verbatim_bits <<= 3;
			aligned_bits = lzx_decode(lzx, &lzx->alignoffset);
			if (aligned_bits < 0)
				return aligned_bits;
		}
		else
		{
			verbatim_bits = lzx_getbits(lzx, offset_bits);
			if (verbatim_bits < 0)
				return verbatim_bits;
			aligned_bits = 0;
		}
		match_offset = (lzx_position_base[position_slot] +
			verbatim_bits + aligned_bits - 2);

		/* Update repeated offset list */
		for (i = (LZX_REPEATED_OFFSETS - 1); i > 0; i--)
			lzx->repeated_offset[i] = lzx->repeated_offset[i - 1];
		lzx->repeated_offset[0] = match_offset;
	}

	/* Copy data */
	if (match_offset > lzx->output.offset)
	{
		return -1;
	}
	if (lzx->output.data)
	{
		copy = &lzx->output.data[lzx->output.offset];
		for (i = 0; i < match_length; i++)
			copy[i] = copy[i - match_offset];
	}
	lzx->output.offset += match_length;

	return 0;
}

/**
 * Translate E8 jump addresses
 *
 * @v lzx    Decompressor
 */
static void lzx_translate_jumps(struct lzx* lzx)
{
	grub_size_t offset;
	grub_int32_t* target;

	/* Sanity check */
	if (lzx->output.offset < 10)
		return;

	/* Scan for jump instructions */
	for (offset = 0; offset < (lzx->output.offset - 10); offset++)
	{
		/* Check for jump instruction */
		if (lzx->output.data[offset] != 0xe8)
			continue;

		/* Translate jump target */
		target = ((grub_int32_t*)&lzx->output.data[offset + 1]);
		if (*target >= 0)
		{
			if (*target < LZX_WIM_MAGIC_FILESIZE)
				*target -= offset;
		}
		else
		{
			if (*target >= -((grub_int32_t)offset))
				*target += LZX_WIM_MAGIC_FILESIZE;
		}
		offset += sizeof(*target);
	}
}

/**
 * Decompress LZX-compressed data
 *
 * @v data    Compressed data
 * @v len    Length of compressed data
 * @v buf    Decompression buffer, or NULL
 * @ret out_len    Length of decompressed data, or negative error
 */
grub_ssize_t
grub_lzx_decompress(const void* data, grub_size_t len, void* buf)
{
	struct lzx lzx;
	unsigned int i;
	int rc;

	/* Sanity check */
	if (len % 2)
	{
		return -1;
	}

	/* Initialise global state, if required */
	if (!lzx_position_base[LZX_POSITION_SLOTS - 1])
	{
		for (i = 1; i < LZX_POSITION_SLOTS; i++)
		{
			lzx_position_base[i] = lzx_position_base[i - 1]
				+ (1 << lzx_footer_bits(i - 1));
		}
	}

	/* Initialise decompressor */
	grub_memset(&lzx, 0, sizeof(lzx));
	lzx.input.data = data;
	lzx.input.len = len;
	lzx.output.data = buf;
	for (i = 0; i < LZX_REPEATED_OFFSETS; i++)
		lzx.repeated_offset[i] = 1;

	/* Process blocks */
	while (lzx.input.offset < lzx.input.len)
	{
		/* Process block header */
		if ((rc = lzx_block_header(&lzx)) != 0)
			return rc;

		/* Process block contents */
		if (lzx.block_type == LZX_BLOCK_UNCOMPRESSED)
		{
			/* Copy uncompressed data */
			if ((rc = lzx_uncompressed(&lzx)) != 0)
				return rc;
		}
		else
		{
			/* Process token stream */
			while (lzx.output.offset < lzx.output.threshold)
			{
				if ((rc = lzx_token(&lzx)) != 0)
					return rc;
			}
		}
	}

	/* Postprocess to undo E8 jump compression */
	if (lzx.output.data)
		lzx_translate_jumps(&lzx);

	return lzx.output.offset;
}
