Universal

main.v

@@ -25,6 +25,12 @@ fn write_lsb(mut b Bits, value u8, length u8) {
 	}
 }
 
+fn write_lsb16(mut b Bits, value u16, length u8) {
+	for i in 0..length {
+		write(mut b, u8((value >> i) & 1))
+	}
+}
+
 fn write_msb(mut b Bits, code u16, length u16) {
 	mut i := int(length) - 1
 	for i >= 0 {
@@ -73,63 +79,132 @@ fn write_deflate_header(mut b Bits) {
 	write_lsb(mut b, 0x01, 2)
 }
 
-fn write_foferk(mut b Bits) {
-	write_msb(mut b, 0x96, 8)
-	write_msb(mut b, 0x9f, 8)
-	write_msb(mut b, 0x96, 8)
-	write_msb(mut b, 0x95, 8)
-	write_msb(mut b, 0xA2, 8)
-	write_msb(mut b, 0x9b, 8)
-	write_msb(mut b, 0x50, 8)
+fn write_text(mut b Bits, text string) {
+	for byte in text.bytes() {
+		if byte <= 143 {
+			write_msb(mut b, u16(byte) + 48, 8)
+		} else {
+			write_msb(mut b, u16(byte) + 256, 9)
 		}
-
-fn write_tokens(mut b Bits, repeat int) {
-	mut remaining := (repeat - 1) * 7
-
-	for remaining >= 252 {
-		write_msb(mut b, 0xC4, 8)
-		write_lsb(mut b, 25, 5)
-		write_msb(mut b, 0x05, 5)
-		write_lsb(mut b, 0x00, 1)
-		remaining -= 252
-	}
-
-	for remaining >= 7 {
-		write_msb(mut b, 0x05, 7)
-		write_msb(mut b, 0x05, 5)
-		write_lsb(mut b, 0x00, 1)
-		remaining -= 7
-	}
-
-	if remaining == 6 {
-		write_msb(mut b, 0x04, 7)
-		write_msb(mut b, 0x05, 5)
-		write_lsb(mut b, 0x00, 1)
-	} else if remaining == 5 {
-		write_msb(mut b, 0x03, 7)
-		write_msb(mut b, 0x05, 5)
-		write_lsb(mut b, 0x00, 1)
-	} else if remaining == 4 {
-		write_msb(mut b, 0x02, 7)
-		write_msb(mut b, 0x05, 5)
-		write_lsb(mut b, 0x00, 1)
-	} else if remaining == 3 {
-		write_msb(mut b, 0x01, 7)
-		write_msb(mut b, 0x05, 5)
-		write_lsb(mut b, 0x00, 1)
 	}
 }
 
-fn write_footer(mut b Bits, repeat int) {
+fn get_distance_code(dist int) (u16, u8, u16) {
+	if dist <= 4 {
+		return u16(dist - 1), 0, 0
+	}
+	
+	mut c := u16(4)
+	mut extra_bits := u8(1)
+	mut base := 5
+
+	for {
+		if dist < base + (1 << extra_bits) {
+			return c, extra_bits, u16(dist - base)
+		}
+	
+		base += (1 << extra_bits)
+		c++
+	
+		if dist < base + (1 << extra_bits) {
+			return c, extra_bits, u16(dist - base)
+		}
+	
+		base += (1 << extra_bits)
+		c++
+		extra_bits++
+	}
+
+	return 0, 0, 0
+}
+
+fn get_length_code(length int) (u16, u8, u16) {
+	if length == 258 {
+		return 285, 0, 0
+	}
+
+	if length <= 10 {
+		return u16(257 + length - 3), 0, 0
+	}
+
+	mut c := u16(265)
+	mut extra_bits := u8(1)
+	mut base := 11
+
+	for {
+		for _ in 0..4 {
+			if length < base + (1 << extra_bits) {
+				return c, extra_bits, u16(length - base)
+			}
+	
+			base += (1 << extra_bits)
+			c++
+	
+			if c >= 285 {
+				return 285, 0, 0
+			}
+		}
+	
+		extra_bits++
+	}
+
+	return 285, 0, 0
+}
+
+fn write_tokens(mut b Bits, repeat int, text string) {
+	if repeat <= 1 || text.len == 0 {
+		return
+	}
+
+	if text.len > 32768 {
+		for _ in 0..(repeat - 1) {
+			write_text(mut b, text)
+		}
+	
+		return
+	}
+
+	mut remaining := (repeat - 1) * text.len
+	dist_code, dist_extra_bits, dist_extra_val := get_distance_code(text.len)
+
+	for remaining > 0 {
+		mut length := 258
+		if remaining < 258 {
+			length = remaining
+		}
+
+		len_code, len_extra_bits, len_extra_val := get_length_code(length)
+
+		if len_code <= 279 {
+			write_msb(mut b, len_code - 256, 7)
+		} else {
+			write_msb(mut b, len_code - 280 + 192, 8)
+		}
+
+		if len_extra_bits > 0 {
+			write_lsb16(mut b, len_extra_val, len_extra_bits)
+		}
+
+		write_msb(mut b, dist_code, 5)
+
+		if dist_extra_bits > 0 {
+			write_lsb16(mut b, dist_extra_val, dist_extra_bits)
+		}
+
+		remaining -= length
+	}
+}
+
+fn write_footer(mut b Bits, repeat int, text string) {
 	write_msb(mut b, 0x00, 7)
 	flush(mut b)
-	write_u32(mut b, crc32(repeat))
-	write_u32(mut b, u32(repeat * 7))
+	write_u32(mut b, crc32(repeat, text))
+	write_u32(mut b, u32(repeat * text.len))
 }
 
 // CRC32
 
-fn crc32(repeat int) u32 {
+fn crc32(repeat int, text string) u32 {
 	mut table := []u32{len: 256}
 	for i in 0..256 {
 		mut c := u32(i)
@@ -145,16 +220,89 @@ fn crc32(repeat int) u32 {
 	}
 
 	mut crc := u32(0xffffffff)
+	bytes := text.bytes()
 	for _ in 0..repeat {
-		for byte in "foferk ".bytes() {
+		for byte in bytes {
 			crc = table[(crc ^ u32(byte)) & 0xFF] ^ (crc >> 8)
 		}
 	}
 	return crc ^ 0xffffffff
 }
 
+// STDIN
+
+fn read_stdin() string {
+	mut lines := ''
+	for {
+		line := os.get_line()
+		if line.len == 0 && os.stdin().eof() {
+			break
+		}
+		lines += line + '\n'
+	}
+	return lines
+}
+
+// ARGS
+
+fn parse_args(text string, repeat int, file string) (string, int, string) {
+	mut text_ := text
+	mut repeat_ := repeat
+	mut file_ := file
+
+	args := os.args[1..]
+
+	if '-h' in args || '--help' in args {
+		println("Usage: foferk [options]")
+		println("")
+		println("Options:")
+		println("  -t, --text <text>")
+		println("  -r, --repeat <count>")
+		println("  -f, --file <path>")
+		println("  -h, --help")
+		println("")
+		exit(0)
+	}
+
+	mut next := -1
+
+	for arg in args {
+		if arg == '-t' || arg == '--text' {
+			next = 0
+		} else if arg == '-r' || arg == '--repeat' {
+			next = 1
+		} else if arg == '-f' || arg == '--file' {
+			next = 2
+		} else {
+			if next == 0 {
+				text_ = arg
+				next = -1
+			} else if next == 1 {
+				repeat_ = arg.int()
+				next = -1
+			} else if next == 2 {
+				file_ = arg
+				next = -1
+			}
+		}
+	}
+
+	return text_, repeat_, file_
+}
+
 fn main() {
-	repeat := 1_000_000
+	mut text := 'foferk'
+	mut repeat := 1_000
+	mut file := 'foferk.gz'
+
+	if os.is_atty(0) == 0 {
+		text = read_stdin()
+	}
+
+	text, repeat, file = parse_args(text, repeat, file)
+
+	text += ' '
+
 	mut b := Bits {
 		bytes: []u8{}
 		current: 0
@@ -163,9 +311,9 @@ fn main() {
 
 	write_gzip_header(mut b)
 	write_deflate_header(mut b)
-	write_foferk(mut b)
-	write_tokens(mut b, repeat)
-	write_footer(mut b, repeat)
+	write_text(mut b, text)
+	write_tokens(mut b, repeat, text)
+	write_footer(mut b, repeat, text)
 
-	os.write_bytes("foferk.gz", b.bytes)!
+	os.write_bytes(file, b.bytes)!
 }