From 83dfd6288b7ab294aa4bd3f11467b3ddffc8854c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Mon, 26 Feb 2024 04:28:37 +0900
Subject: [PATCH] feat: add writer

---
 README.md      |  1 -
 reader.go      | 10 ++++++++
 spb.go         | 35 ----------------------------
 utils.go       | 36 +++++++++++++++++++++++++++++
 writer.go      | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++
 writer_test.go | 36 +++++++++++++++++++++++++++++
 6 files changed, 144 insertions(+), 36 deletions(-)
 create mode 100644 utils.go
 create mode 100644 writer.go
 create mode 100644 writer_test.go

diff --git a/README.md b/README.md
index 9c52639..acabf76 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,3 @@
 # go-simple-protobuf
-> Note: Only supports reading.
 
 Golang implementation of [simple-protobuf](https://github.com/fumiama/simple-protobuf)
diff --git a/reader.go b/reader.go
index c8aa5db..9363639 100644
--- a/reader.go
+++ b/reader.go
@@ -65,3 +65,13 @@ func (it *Iterator) Bytes() []byte {
 func (it *Iterator) String() string {
 	return string(it.v)
 }
+
+// Cap structlen 结构体带 padding 的长度
+func (it *Iterator) Cap() uint32 {
+	return it.structlen
+}
+
+// Len reallen (迭代器目前迭代到的实值总长, 不带 padding)
+func (it *Iterator) Len() uint32 {
+	return it.reallen
+}
diff --git a/spb.go b/spb.go
index 943feb3..a8e9d44 100644
--- a/spb.go
+++ b/spb.go
@@ -2,41 +2,6 @@ package spb
 
 import "io"
 
-func ReadNum(r io.Reader) (n uint32, cnt uint32, err error) {
-	var buf [1]byte
-	for cnt < 5 {
-		_, err = r.Read(buf[:])
-		if err != nil {
-			return
-		}
-		n |= uint32(buf[0]&0x7f) << (7 * cnt)
-		cnt++
-		if buf[0]&0x80 == 0 {
-			break
-		}
-	}
-	return
-}
-
-/*
-func WriteNum(w io.Writer, n uint32) (cnt int, err error) {
-	var buf [1]byte
-	for n > 0 {
-		buf[0] = uint8(n & 0x7f)
-		if n>>7 > 0 {
-			buf[0] |= 0x80
-		}
-		_, err = w.Write(buf[:])
-		if err != nil {
-			return
-		}
-		n >>= 7
-		cnt++
-	}
-	return
-}
-*/
-
 type SimplePB struct {
 	StructLen, RealLen uint32
 	Target             [][]byte
diff --git a/utils.go b/utils.go
new file mode 100644
index 0000000..a5a14f9
--- /dev/null
+++ b/utils.go
@@ -0,0 +1,36 @@
+package spb
+
+import "io"
+
+func ReadNum(r io.Reader) (n uint32, cnt uint32, err error) {
+	var buf [1]byte
+	for cnt < 5 {
+		_, err = r.Read(buf[:])
+		if err != nil {
+			return
+		}
+		n |= uint32(buf[0]&0x7f) << (7 * cnt)
+		cnt++
+		if buf[0]&0x80 == 0 {
+			break
+		}
+	}
+	return
+}
+
+func WriteNum(w io.Writer, n uint32) (cnt uint32, err error) {
+	var buf [1]byte
+	for n > 0 {
+		buf[0] = uint8(n & 0x7f)
+		if n>>7 > 0 {
+			buf[0] |= 0x80
+		}
+		_, err = w.Write(buf[:])
+		if err != nil {
+			return
+		}
+		n >>= 7
+		cnt++
+	}
+	return
+}
diff --git a/writer.go b/writer.go
new file mode 100644
index 0000000..4f061bf
--- /dev/null
+++ b/writer.go
@@ -0,0 +1,62 @@
+package spb
+
+import (
+	"bytes"
+	"io"
+)
+
+type Constructor struct {
+	structlen, reallen uint32
+	body               bytes.Buffer
+}
+
+// NewWriter 新的写入构造器
+func NewWriter(buf []byte) (c Constructor) {
+	c.body = *bytes.NewBuffer(buf)
+	return
+}
+
+// WriteString 按字符串写入一项
+func (c *Constructor) WriteString(s string, cap uint32) error {
+	if int(cap) < len(s) {
+		return ErrInvalidDataLen
+	}
+	cnt, err := WriteNum(&c.body, cap)
+	if err != nil {
+		return err
+	}
+	c.structlen += cap
+	c.reallen += cnt
+	cnt, err = WriteNum(&c.body, uint32(len(s)))
+	if err != nil {
+		return err
+	}
+	c.reallen += cnt
+	_, err = c.body.WriteString(s)
+	if err != nil {
+		return err
+	}
+	c.reallen += uint32(len(s))
+	return nil
+}
+
+// Cap structlen 目前写入的总长, 带 padding
+func (c *Constructor) Cap() uint32 {
+	return c.structlen
+}
+
+// Len reallen 目前写入的实值总长, 不带 padding
+func (c *Constructor) Len() uint32 {
+	return c.reallen
+}
+
+// WriteTo 将当前结果写出, 返回写入的实际长度
+func (c *Constructor) WriteTo(w io.Writer) (int64, error) {
+	cnt, err := WriteNum(w, c.structlen)
+	if err != nil {
+		return int64(cnt), err
+	}
+	n, err := w.Write(c.body.Bytes())
+	cnt += uint32(n)
+	return int64(cnt), err
+}
diff --git a/writer_test.go b/writer_test.go
new file mode 100644
index 0000000..cc59ef9
--- /dev/null
+++ b/writer_test.go
@@ -0,0 +1,36 @@
+package spb
+
+import (
+	"bufio"
+	"bytes"
+	"os"
+	"strings"
+	"testing"
+)
+
+func TestWriter(t *testing.T) {
+	f := bytes.NewBuffer(make([]byte, 0, 65536))
+	ft, err := os.Open("dict.txt")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ft.Close()
+	sc := bufio.NewScanner(ft)
+	for sc.Scan() {
+		c := Constructor{}
+		for _, s := range strings.Split(sc.Text(), "\t") {
+			err = c.WriteString(s, 127)
+			if err != nil {
+				t.Fatal(err)
+			}
+		}
+		c.WriteTo(f)
+	}
+	real, err := os.ReadFile("dict.sp")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(real, f.Bytes()) {
+		t.Fail()
+	}
+}