means (#3)

Co-authored-by: Evan Burkey <dev@fputs.com>
Reviewed-on: https://git.fputs.com/fputs/protohackers/pulls/3

cmd/meanstoanend/main.go

@@ -0,0 +1,60 @@
+package main
+
+import (
+	"encoding/binary"
+	"io"
+	"log"
+	"net"
+
+	"protohackers/pkg/conn"
+)
+
+func main() {
+	err := conn.StartSimple(meansToAnEnd)
+	if err != nil {
+		log.Fatalln(err)
+	}
+}
+
+func meansToAnEnd(c net.Conn) {
+	defer c.Close()
+	var ledger = make(map[int32]int32)
+	buf := make([]byte, 9)
+
+	for {
+		if _, err := io.ReadFull(c, buf); err != nil {
+			return
+		}
+		if _, err := c.Write(process(buf, ledger)); err != nil {
+			return
+		}
+	}
+}
+
+func process(buf []byte, ledger map[int32]int32) []byte {
+	i1, i2 := int32(binary.BigEndian.Uint32(buf[1:5])), int32(binary.BigEndian.Uint32(buf[5:]))
+
+	switch buf[0] {
+	case 'I':
+		ledger[i1] = i2
+	case 'Q':
+		var c int64
+		var t int64
+		for k, v := range ledger {
+			if i1 <= k && k <= i2 {
+				t += int64(v)
+				c++
+			}
+		}
+		if c == 0 {
+			c = 1 /* Can't divide by zero! */
+		}
+		mean := t / c
+		r := make([]byte, 4)
+		binary.BigEndian.PutUint32(r, uint32(mean))
+		return r
+	default:
+		return nil
+	}
+	return nil
+}

cmd/primetime/main.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"bufio"
 	"encoding/json"
 	"log"
 	"math/big"
@@ -29,13 +28,7 @@ func main() {
 }
 
 func primetime(c net.Conn) {
-	defer c.Close()
+	if err := conn.ReadAndRespond(c, process); err != nil {
-	sc := bufio.NewScanner(c)
-	for sc.Scan() {
-		response := process(sc.Bytes())
-		c.Write(response)
-	}
-	if err := sc.Err(); err != nil {
 		log.Fatalln(err)
 	}
 }

pkg/conn/conn.go

@@ -1,6 +1,7 @@
 package conn
 
 import (
+	"bufio"
 	"fmt"
 	"net"
 )
@@ -30,3 +31,19 @@ func StartSimple(handler func(conn net.Conn)) error {
 		}
 	}
 }
+
+func ReadAndRespond(c net.Conn, process func([]byte) []byte) error {
+	defer c.Close()
+	sc := bufio.NewScanner(c)
+	for sc.Scan() {
+		response := process(sc.Bytes())
+		_, err := c.Write(response)
+		if err != nil {
+			return err
+		}
+	}
+	if err := sc.Err(); err != nil {
+		return err
+	}
+	return nil
+}