From 14dd858b0439e11bd2cf849371e64debe03e9e90 Mon Sep 17 00:00:00 2001
From: wenyekui <wyk4true@gmail.com>
Date: Sat, 16 Aug 2014 13:43:21 +0800
Subject: [PATCH] bug fix: set zset size

---
 ledis/scan_test.go   | 22 ----------------------
 ledis/t_zset.go      | 29 ++++++++++-------------------
 ledis/t_zset_test.go | 23 ++++++++++++++++++++++-
 3 files changed, 32 insertions(+), 42 deletions(-)

diff --git a/ledis/scan_test.go b/ledis/scan_test.go
index 1e14a60..c2dbd34 100644
--- a/ledis/scan_test.go
+++ b/ledis/scan_test.go
@@ -1,8 +1,6 @@
 package ledis
 
 import (
-	"fmt"
-	"github.com/siddontang/ledisdb/store"
 	"testing"
 )
 
@@ -198,26 +196,6 @@ func TestDBBScan(t *testing.T) {
 		t.Fatal(err.Error())
 	}
 
-	ek1 := db.bEncodeMetaKey(k1)
-	fmt.Printf("%x\n", ek1)
-	ek2 := db.bEncodeMetaKey(k2)
-	fmt.Printf("%x\n", ek2)
-	ek3 := db.bEncodeMetaKey(k3)
-	fmt.Printf("%x\n", ek3)
-
-	start := db.bEncodeMetaKey(nil)
-	fmt.Printf("start: %x\n", start)
-	end := db.bEncodeMetaKey(nil)
-	end[len(end)-1] = BitMetaType + 1
-	fmt.Printf("end: %x\n", end)
-
-	it := db.db.RangeLimitIterator(start, end, store.RangeClose, 0, 4)
-
-	for ; it.Valid(); it.Next() {
-		fmt.Printf("%x\n", it.RawKey())
-	}
-	it.Close()
-
 	if v, err := db.BScan(nil, 1, true); err != nil {
 		t.Fatal(err)
 	} else if len(v) != 1 {
diff --git a/ledis/t_zset.go b/ledis/t_zset.go
index bed0c6f..5458a01 100644
--- a/ledis/t_zset.go
+++ b/ledis/t_zset.go
@@ -873,29 +873,25 @@ func (db *DB) ZUnionStore(destKey []byte, srcKeys [][]byte, weights []int64, agg
 
 	db.zDelete(t, destKey)
 
-	var num int64 = 0
 	for member, score := range destMap {
 		if err := checkZSetKMSize(destKey, []byte(member)); err != nil {
 			return 0, err
 		}
 
-		if n, err := db.zSetItem(t, destKey, score, []byte(member)); err != nil {
+		if _, err := db.zSetItem(t, destKey, score, []byte(member)); err != nil {
 			return 0, err
-		} else if n == 0 {
-			//add new
-			num++
 		}
 	}
 
-	if _, err := db.zIncrSize(t, destKey, num); err != nil {
-		return 0, err
-	}
+	var num = int64(len(destMap))
+	sk := db.zEncodeSizeKey(destKey)
+	t.Put(sk, PutInt64(num))
 
 	//todo add binlog
 	if err := t.Commit(); err != nil {
 		return 0, err
 	}
-	return int64(len(destMap)), nil
+	return num, nil
 }
 
 func (db *DB) ZInterStore(destKey []byte, srcKeys [][]byte, weights []int64, aggregate byte) (int64, error) {
@@ -947,28 +943,23 @@ func (db *DB) ZInterStore(destKey []byte, srcKeys [][]byte, weights []int64, agg
 
 	db.zDelete(t, destKey)
 
-	var num int64 = 0
 	for member, score := range destMap {
 		if err := checkZSetKMSize(destKey, []byte(member)); err != nil {
 			return 0, err
 		}
-
-		if n, err := db.zSetItem(t, destKey, score, []byte(member)); err != nil {
+		if _, err := db.zSetItem(t, destKey, score, []byte(member)); err != nil {
 			return 0, err
-		} else if n == 0 {
-			//add new
-			num++
 		}
 	}
 
-	if _, err := db.zIncrSize(t, destKey, num); err != nil {
-		return 0, err
-	}
+	var num int64 = int64(len(destMap))
+	sk := db.zEncodeSizeKey(destKey)
+	t.Put(sk, PutInt64(num))
 	//todo add binlog
 	if err := t.Commit(); err != nil {
 		return 0, err
 	}
-	return int64(len(destMap)), nil
+	return num, nil
 }
 
 func (db *DB) ZScan(key []byte, count int, inclusive bool) ([][]byte, error) {
diff --git a/ledis/t_zset_test.go b/ledis/t_zset_test.go
index 10782d1..212a02d 100644
--- a/ledis/t_zset_test.go
+++ b/ledis/t_zset_test.go
@@ -253,6 +253,9 @@ func TestZUnionStore(t *testing.T) {
 	weights := []int64{1, 2}
 
 	out := []byte("out")
+
+	db.ZAdd(out, ScorePair{3, []byte("out")})
+
 	n, err := db.ZUnionStore(out, keys, weights, AggregateSum)
 	if err != nil {
 		t.Fatal(err.Error())
@@ -296,6 +299,15 @@ func TestZUnionStore(t *testing.T) {
 	if n != 3 {
 		t.Fatal("invalid value ", v)
 	}
+
+	n, err = db.ZCard(out)
+
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+	if n != 3 {
+		t.Fatal("invalid value ", n)
+	}
 }
 
 func TestZInterStore(t *testing.T) {
@@ -314,6 +326,8 @@ func TestZInterStore(t *testing.T) {
 	weights := []int64{2, 3}
 	out := []byte("out")
 
+	db.ZAdd(out, ScorePair{3, []byte("out")})
+
 	n, err := db.ZInterStore(out, keys, weights, AggregateSum)
 	if err != nil {
 		t.Fatal(err.Error())
@@ -329,7 +343,6 @@ func TestZInterStore(t *testing.T) {
 		t.Fatal("invalid value ", v)
 	}
 
-	out = []byte("out")
 	n, err = db.ZInterStore(out, keys, weights, AggregateMin)
 	if err != nil {
 		t.Fatal(err.Error())
@@ -355,4 +368,12 @@ func TestZInterStore(t *testing.T) {
 		t.Fatal("invalid value ", n)
 	}
 
+	n, err = db.ZCard(out)
+
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+	if n != 1 {
+		t.Fatal("invalid value ", n)
+	}
 }