@@ -1,56 +1,50 @@
 package main
 
 import (
 	mtf "0dev.org/encoding/mtf"
 	iou "0dev.org/ioutil"
+	"flag"
 	"fmt"
 	"io"
 	"os"
 )
 
 func main() {
-	var code int
-	switch {
-	case len(os.Args) == 1:
-		code = transform(os.Stdout, os.Stdin)
-	case len(os.Args) == 2 && os.Args[1] == "-d":
-		code = reverse(os.Stdout, os.Stdin)
-	default:
-		fmt.Fprintln(os.Stdout, "Usage: mtf [-d]")
-	}
+	d := flag.Bool("d", false, "Use to toggle decode mode")
+	flag.Parse()
+
+	var code int
+	if *d {
+		code = decode(os.Stdout, os.Stdin)
+	} else {
+		code = encode(os.Stdout, os.Stdin)
+	}
+
 	os.Exit(code)
 }
 
 // Transforms the data according to the move-to-front algorithm
 // I/O is buffered for better performance
-func transform(output io.Writer, input io.Reader) int {
+func encode(output io.Writer, input io.Reader) int {
 	var (
 		err     error
-		buffer  io.Writer = iou.SizedWriter(output, 4096)
-		encoder io.Writer = mtf.Encoder(buffer)
+		encoder io.Reader = mtf.Encoder(iou.SizedReader(input, 4096))
 	)
 
-	_, err = io.Copy(encoder, input)
+	_, err = io.Copy(output, encoder)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, "Error while encoding.\n", err)
 		return 1
 	}
 
-	// Flush the buffer
-	_, err = buffer.Write(nil)
-	if err != nil {
-		fmt.Fprintln(os.Stderr, "Error while flushing output buffer.\n", err)
-		return 1
-	}
-
 	return 0
 }
 
 // Reverses MTF`ed data and writes back the original bytes
 // I/O is buffered for better performance
-func reverse(output io.Writer, input io.Reader) int {
+func decode(output io.Writer, input io.Reader) int {
 	var (
 		err     error
 		decoder io.Reader = mtf.Decoder(iou.SizedReader(input, 4096))
 	)