new pkg ic

pkg/assert/assert.go

@@ -6,10 +6,7 @@
 //
 // Author: Chef (191201771@qq.com)
 
-// package assert 提供了单元测试时的断言功能，减少一些模板代码
-//
-// 代码参考了 https://github.com/stretchr/testify
-//
+// Package assert 提供了单元测试时的断言功能，减少一些模板代码
 package assert
 
 import (
@@ -17,6 +14,7 @@ import (
 	"reflect"
 )
 
+// 单元测试中的 *testing.T 和 *testing.B 都满足该接口
 type TestingT interface {
 	Errorf(format string, args ...interface{})
 }

pkg/assert/assert_test.go

@@ -6,10 +6,12 @@
 //
 // Author: Chef (191201771@qq.com)
 
-package assert
+package assert_test
 
 import (
+	"errors"
 	"fmt"
+	"github.com/q191201771/naza/pkg/assert"
 	"testing"
 )
 
@@ -23,39 +25,35 @@ func (mtt MockTestingT) Errorf(format string, args ...interface{}) {
 }
 func TestEqual(t *testing.T) {
 	// 测试Equal
-	Equal(t, nil, nil)
-	Equal(t, nil, nil, "fxxk.")
-	Equal(t, 1, 1)
-	Equal(t, "aaa", "aaa")
+	assert.Equal(t, nil, nil)
+	assert.Equal(t, nil, nil, "fxxk.")
+	assert.Equal(t, 1, 1)
+	assert.Equal(t, "aaa", "aaa")
 	var ch chan struct{}
-	Equal(t, nil, ch)
+	assert.Equal(t, nil, ch)
 	var m map[string]string
-	Equal(t, nil, m)
+	assert.Equal(t, nil, m)
 	var p *int
-	Equal(t, nil, p)
+	assert.Equal(t, nil, p)
 	var i interface{}
-	Equal(t, nil, i)
+	assert.Equal(t, nil, i)
 	var b []byte
-	Equal(t, nil, b)
-	//Equal(t, nil, errors.New("mock error"))
+	assert.Equal(t, nil, b)
 
-	// 测试isNil
-	Equal(t, true, isNil(nil))
-	Equal(t, false, isNil("aaa"))
-	// 测试equal
-	Equal(t, false, equal([]byte{}, "aaa"))
-	Equal(t, true, equal([]byte{}, []byte{}))
-	Equal(t, true, equal([]byte{0, 1, 2}, []byte{0, 1, 2}))
+	assert.Equal(t, []byte{}, []byte{})
+	assert.Equal(t, []byte{0, 1, 2}, []byte{0, 1, 2})
 
 	// 测试Equal失败
 	var mtt MockTestingT
-	Equal(mtt, false, isNil(nil))
+	assert.Equal(mtt, nil, 1)
+	assert.Equal(mtt, []byte{}, "aaa")
+	assert.Equal(mtt, nil, errors.New("mock error"))
 }
 
 func TestIsNotNil(t *testing.T) {
-	IsNotNil(t, 1)
+	assert.IsNotNil(t, 1)
 
 	// 测试IsNotNil失败
 	var mtt MockTestingT
-	IsNotNil(mtt, nil)
+	assert.IsNotNil(mtt, nil)
 }

pkg/ic/ic_test.go

@@ -0,0 +1,115 @@
+package ic
+
+import (
+	"bytes"
+	"io/ioutil"
+	"log"
+	"os"
+	"strconv"
+	"testing"
+)
+
+// 从文件中读取 uid 列表
+func obtainUIDList(filename string) (uids IDSlice) {
+	fp, err := os.Open(filename)
+	if err != nil {
+		panic(err)
+	}
+	buf, err := ioutil.ReadAll(fp)
+	if err != nil {
+		panic(err)
+	}
+	lines := bytes.Split(buf, []byte("\n"))
+	for _, line := range lines {
+		if len(line) == 0 {
+			continue
+		}
+		item, err := strconv.ParseUint(string(line), 10, 32)
+		if err != nil {
+			panic(err)
+		}
+		uids = append(uids, uint32(item))
+	}
+	return uids
+}
+
+//var FILENAME = "uid.txt"
+
+func marshalWrap(ids IDSlice) (ret []byte) {
+	log.Println("> sort.")
+	sortIDSlice(ids)
+	log.Println("< sort.")
+
+	log.Println("> marshal.")
+	//var oc OriginCompressor
+	//ret = oc.Marshal(ids)
+
+	var lfc LFCompressor
+	lfc.FB = 4
+	ret = lfc.Marshal(ids)
+	log.Println("< marshal.")
+
+	log.Println("> zlib. len:", len(ret))
+	ret = zlibWrite(ret)
+	log.Println("< zlib. len:", len(ret))
+	return
+}
+
+func unmarshalWrap(b []byte) (ret IDSlice) {
+	b = zlibRead(b)
+
+	//var oc OriginCompressor
+	//ret = oc.Unmarshal(b)
+
+	var lfc LFCompressor
+	lfc.FB = 4
+	ret = lfc.Unmarshal(b)
+	return
+}
+
+func TestIC(t *testing.T) {
+	log.SetFlags(log.Lmicroseconds)
+
+	// 单元测试 case
+	uidss := []IDSlice{
+		{1, 2, 3, 18, 32, 100},
+		{1, 2, 3, 18, 32},
+		{1, 2, 3, 18},
+		{1, 2, 3, 17},
+		{1, 2, 3, 16},
+		{1, 2, 3, 15, 16, 17, 18},
+		{1, 2, 3, 15, 16, 17},
+		{1, 2, 3, 15, 16},
+		{1, 2, 3, 15},
+		{1, 2, 3},
+		{1, 2},
+		{1},
+	}
+
+	// 从文件加载 uid 白名单
+	//uids := obtainUIDList(FILENAME)
+	//var uidss []IDSlice
+	//uidss = append(uidss, uids)
+
+	for _, uids := range uidss {
+		log.Println("-----")
+		log.Println("in uid len:", len(uids))
+
+		b := marshalWrap(uids)
+		log.Println("len(b):", len(b))
+
+		uids2 := unmarshalWrap(b)
+		log.Println("out uid len:", len(uids2))
+
+		// assert check
+		if len(uids) != len(uids2) {
+			panic(0)
+		}
+		for i := range uids {
+			if uids[i] != uids2[i] {
+				panic(0)
+			}
+		}
+		log.Println("-----")
+	}
+}

pkg/ic/lf_compressor.go

@@ -0,0 +1,102 @@
+package ic
+
+import (
+	"encoding/binary"
+)
+
+type LFCompressor struct {
+	FB uint32           // 用几个字节的 bit 表示跟随的数据
+	oc OriginCompressor // FB 为0时，退化成使用 OriginCompressor
+}
+
+func (lfc *LFCompressor) Marshal(ids IDSlice) (ret []byte) {
+	if lfc.FB == 0 {
+		return lfc.oc.Marshal(ids)
+	}
+
+	lBuf := make([]byte, 4)
+	fBuf := make([]byte, lfc.FB)
+
+	maxDiff := 8 * lfc.FB
+
+	var hasLeader bool
+	var leader uint32
+	var stage int
+	for i := range ids {
+		if !hasLeader {
+			stage = 1
+			leader = ids[i]
+			hasLeader = true
+			continue
+		}
+
+		diff := uint32(ids[i] - leader)
+
+		if diff > maxDiff {
+			binary.LittleEndian.PutUint32(lBuf, leader)
+			ret = append(ret, lBuf...)
+			ret = append(ret, fBuf...)
+
+			resetBuf(fBuf)
+			stage = 2
+			leader = ids[i]
+		} else {
+			stage = 3
+			fBuf[(diff-1)/8] = fBuf[(diff-1)/8] | (1 << byte((diff-1)%8))
+		}
+	}
+
+	switch stage {
+	case 1:
+		binary.LittleEndian.PutUint32(lBuf, leader)
+		ret = append(ret, lBuf...)
+		dummy := make([]byte, lfc.FB)
+		ret = append(ret, dummy...)
+	case 2:
+		binary.LittleEndian.PutUint32(lBuf, leader)
+		ret = append(ret, lBuf...)
+		dummy := make([]byte, lfc.FB)
+		ret = append(ret, dummy...)
+	case 3:
+		binary.LittleEndian.PutUint32(lBuf, leader)
+		ret = append(ret, lBuf...)
+		ret = append(ret, fBuf...)
+	}
+	return
+}
+
+func (lfc *LFCompressor) Unmarshal(b []byte) (ids IDSlice) {
+	if lfc.FB == 0 {
+		return lfc.oc.Unmarshal(b)
+	}
+
+	isLeaderStage := true
+	var item uint32
+	var leader uint32
+	var index uint32
+	for {
+		if isLeaderStage {
+			leader = binary.LittleEndian.Uint32(b[index:])
+			ids = append(ids, leader)
+			isLeaderStage = false
+			index += 4
+		} else {
+			for i := uint32(0); i < lfc.FB; i++ {
+				for j := uint32(0); j < 8; j++ {
+					if ((b[index+i] >> j) & 1) == 1 {
+						item = leader + (i * 8) + j + 1
+						ids = append(ids, item)
+					}
+				}
+			}
+
+			isLeaderStage = true
+			index += lfc.FB
+		}
+
+		if int(index) == len(b) {
+			break
+		}
+	}
+	return
+}

pkg/ic/origin_compressor.go

@@ -0,0 +1,23 @@
+package ic
+
+import "encoding/binary"
+
+type OriginCompressor struct {
+}
+
+func (oc *OriginCompressor) Marshal(ids IDSlice) (ret []byte) {
+	ret = make([]byte, len(ids)*4)
+	for i, id := range ids {
+		binary.LittleEndian.PutUint32(ret[i*4:], id)
+	}
+	return
+}
+
+func (oc *OriginCompressor) Unmarshal(b []byte) (ids IDSlice) {
+	n := len(b) / 4
+	for i := 0; i < n; i++ {
+		id := binary.LittleEndian.Uint32(b[i*4:])
+		ids = append(ids, id)
+	}
+	return
+}

pkg/ic/util.go

@@ -0,0 +1,58 @@
+package ic
+
+import (
+	"bytes"
+	"compress/zlib"
+	"io/ioutil"
+	"sort"
+)
+
+type IDSlice []uint32
+
+func (a IDSlice) Len() int           { return len(a) }
+func (a IDSlice) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
+func (a IDSlice) Less(i, j int) bool { return a[i] < a[j] }
+
+func resetBuf(b []byte) []byte {
+	for i := 0; i < len(b); i++ {
+		b[i] = 0
+	}
+	return b
+}
+
+func sortIDSlice(ids IDSlice) {
+	sort.Sort(ids)
+}
+
+func zlibWrite(in []byte) []byte {
+	var b bytes.Buffer
+	w := zlib.NewWriter(&b)
+	_, _ = w.Write(in)
+	_ = w.Close()
+	return b.Bytes()
+}
+
+func zlibRead(in []byte) (ret []byte) {
+	b := bytes.NewReader(in)
+	r, _ := zlib.NewReader(b)
+	ret, _ = ioutil.ReadAll(r)
+	return
+}
+
+//func isBufEmpty(b []byte) bool {
+//	for i := 0; i < len(b); i++ {
+//		if b[i] != 0 {
+//			return false
+//		}
+//	}
+//	return true
+//}
+//
+//func dumpIDSlice(ids IDSlice, filename string) {
+//	fp, _ := os.Create(filename)
+//	for _, id := range ids {
+//		_, _ = fp.WriteString(fmt.Sprintf("%d", id))
+//		_, _ = fp.WriteString("\n")
+//	}
+//	_ = fp.Close()
+//}