Add BoltStorage

cmd/gen.go

@@ -45,7 +45,11 @@ var genCmd = &cobra.Command{
 			words = append(words, record[colIndex])
 		}
 
-		iroha := lib.NewIroha(words, &lib.NoStorage{}, &lib.DepthOptions{})
+		boltStorage, err := lib.NewBoltStorage("test.db")
+		if err != nil {
+			panic(err)
+		}
+		iroha := lib.NewIroha(words, boltStorage, config.DepthOptions)
 		iroha.PrintWordCountMap()
 		iroha.PrintWordByKatakanaMap()
 		rowIndicesList, err := iroha.Search()
@@ -127,6 +131,9 @@ func init() {
 	if err := registerIntToFlags(genCmd, flagKeys.MaxLogDepth, 0, "max log depth"); err != nil {
 		panic(err)
 	}
+	if err := registerIntToFlags(genCmd, flagKeys.MaxStorageDepth, -1, "max storage depth"); err != nil {
+		panic(err)
+	}
 }
 
 func registerIntToFlags(cmd *cobra.Command, name string, value int, usage string) error {

gen/gen.go

@@ -21,6 +21,7 @@ type FlagKeys struct {
 	MinParallelDepth string
 	MaxLogDepth      string
 	MaxParallelDepth string
+	MaxStorageDepth  string
 	OutputMode       string
 }
 
@@ -31,14 +32,15 @@ func NewFlagKeys() *FlagKeys {
 		MinParallelDepth: "min-p-depth",
 		MaxLogDepth:      "max-log-depth",
 		MaxParallelDepth: "max-p-depth",
+		MaxStorageDepth:  "max-s-depth",
 		OutputMode:       "output-mode",
 	}
 }
 
 type Config struct {
 	FilePath     string
 	ColName      string
-	depthOptions *lib.DepthOptions
+	DepthOptions *lib.DepthOptions
 	OutputMode   OutputMode
 }
 
@@ -47,10 +49,11 @@ func NewConfigFromViper() *Config {
 	return &Config{
 		FilePath: viper.GetString(flagKeys.File),
 		ColName:  viper.GetString(flagKeys.ColName),
-		depthOptions: &lib.DepthOptions{
+		DepthOptions: &lib.DepthOptions{
 			MinParallel: viper.GetInt(flagKeys.MinParallelDepth),
 			MaxParallel: viper.GetInt(flagKeys.MaxParallelDepth),
 			MaxLog:      viper.GetInt(flagKeys.MaxLogDepth),
+			MaxStorage:  viper.GetInt(flagKeys.MaxStorageDepth),
 		},
 		OutputMode: OutputMode(viper.GetString(flagKeys.OutputMode)),
 	}

go.mod

@@ -8,6 +8,7 @@ require (
 	github.com/k0kubun/pp v3.0.1+incompatible
 	github.com/mattn/go-colorable v0.1.2 // indirect
 	github.com/mitchellh/go-homedir v1.1.0
+	github.com/mpppk/bbolt v1.3.3
 	github.com/pkg/errors v0.8.1
 	github.com/rhysd/go-github-selfupdate v1.1.0
 	github.com/spf13/cobra v0.0.5

go.sum

@@ -81,6 +81,8 @@ github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/mapstructure v1.1.2 h1:fmNYVwqnSfB9mZU6OS2O6GsXM+wcskZDuKQzvN1EDeE=
 github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
+github.com/mpppk/bbolt v1.3.3 h1:AGktmY9GtlsLvZZuUFIgrjiLjdI/8c1fWorFRgLDyLY=
+github.com/mpppk/bbolt v1.3.3/go.mod h1:mXR1eqEWlzD6jgmHhpK8y3UTgk12JjbfqHigJjJwoKk=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/oklog/ulid v1.3.1/go.mod h1:CirwcVhetQ6Lv90oh/F+FBtV6XMibvdAFo93nm5qn4U=
 github.com/onsi/ginkgo v1.6.0 h1:Ix8l273rp3QzYgXSR+c8d1fTG7UPgYkOSELPhiY/YGw=

lib/iroha.go

@@ -43,7 +43,7 @@ func (i *Iroha) PrintWordByKatakanaMap() {
 func (i *Iroha) Search() (rowIndicesList [][]int, err error) {
 	katakanaBitsAndWordsList := i.katakana.ListSortedKatakanaBitsAndWords()
 	i.log = NewLog(katakanaBitsAndWordsList, i.depths.MaxLog, i.depths.MinParallel)
-	wordsList, _, err := i.searchByBits([]int{}, katakanaBitsAndWordsList, WordBits(0))
+	wordsList, _, _, err := i.searchByBits([]int{}, katakanaBitsAndWordsList, WordBits(0))
 	if err != nil {
 		return nil, err
 	}
@@ -57,27 +57,28 @@ func (i *Iroha) Search() (rowIndicesList [][]int, err error) {
 	return
 }
 
-func (i *Iroha) f(word *Word, usedIndices []int, katakanaBitsAndWords []*KatakanaBitsAndWords, remainKatakanaBits WordBits) ([][]*Word, bool, error) {
+func (i *Iroha) f(word *Word, usedIndices []int, katakanaBitsAndWords []*KatakanaBitsAndWords, remainKatakanaBits WordBits) ([][]*Word, bool, bool, error) {
 	var results [][]*Word
 	if remainKatakanaBits.HasDuplicatedKatakana(word.Bits) {
-		return nil, false, nil
+		return nil, false, false, nil
 	}
 	newRemainKatakanaBits := remainKatakanaBits.Merge(word.Bits)
-	newIrohaWordIdLists, ok, err := i.searchByBits(usedIndices, katakanaBitsAndWords[1:], newRemainKatakanaBits)
+	newIrohaWordIdLists, ok, cacheUsed, err := i.searchByBits(usedIndices, katakanaBitsAndWords[1:], newRemainKatakanaBits)
 	if err != nil {
-		return nil, false, err
+		return nil, false, false, err
 	}
 	if ok {
 		for _, newIrohaWordList := range newIrohaWordIdLists {
 			newIrohaWordList = append(newIrohaWordList, word)
 			results = append(results, newIrohaWordList)
 		}
 	}
-	return results, true, nil
+	return results, true, cacheUsed, nil
 }
 
 func (i *Iroha) gf(word *Word, usedIndices []int, katakanaBitsAndWords []*KatakanaBitsAndWords, remainKatakanaBits WordBits, wordListChan chan<- []*Word) error {
-	wordLists, ok, err := i.f(word, usedIndices, katakanaBitsAndWords, remainKatakanaBits)
+	// FIXME: check cache
+	wordLists, ok, _, err := i.f(word, usedIndices, katakanaBitsAndWords, remainKatakanaBits)
 	if err != nil {
 		return err
 	}
@@ -89,32 +90,33 @@ func (i *Iroha) gf(word *Word, usedIndices []int, katakanaBitsAndWords []*Kataka
 	return nil
 }
 
-func (i *Iroha) searchByBits(usedIndices []int, katakanaBitsAndWords []*KatakanaBitsAndWords, remainKatakanaBits WordBits) ([][]*Word, bool, error) {
+func (i *Iroha) searchByBits(usedIndices []int, katakanaBitsAndWords []*KatakanaBitsAndWords, remainKatakanaBits WordBits) ([][]*Word, bool, bool, error) {
 	remainKatakanaNum := bits.OnesCount64(uint64(remainKatakanaBits))
 	if remainKatakanaNum == int(KatakanaLen) {
-		return [][]*Word{{}}, true, nil
+		return [][]*Word{{}}, true, false, nil
 	}
 
 	if len(katakanaBitsAndWords) == 0 {
-		return nil, false, nil
+		return nil, false, false, nil
 	}
 
 	katakanaAndWordBits := katakanaBitsAndWords[0]
 	if len(katakanaAndWordBits.Words) == 0 {
-		return nil, false, nil
+		return nil, false, false, nil
 	}
 
-	if results, ok, err := i.storage.Get(usedIndices); err != nil {
-		return nil, false, err
-	} else if ok {
-		return results, true, nil
+	depth := int(KatakanaLen) - len(katakanaBitsAndWords)
+
+	if depth <= i.depths.MaxStorage {
+		if results, ok, err := i.storage.Get(usedIndices); err != nil {
+			return nil, false, false, err
+		} else if ok {
+			return results, true, true, nil
+		}
 	}
 
-	depth := int(KatakanaLen) - len(katakanaBitsAndWords)
 	var irohaWordLists [][]*Word
-
 	goroutineMode := depth >= i.depths.MinParallel && depth <= i.depths.MaxParallel
-
 	if goroutineMode {
 		eg := errgroup.Group{}
 		wordListChan := make(chan []*Word, 100)
@@ -148,41 +150,56 @@ func (i *Iroha) searchByBits(usedIndices []int, katakanaBitsAndWords []*Katakana
 				irohaWordLists = append(irohaWordLists, wordList)
 			case err, ok := <-errChan:
 				if !ok {
-					return nil, false, fmt.Errorf("unexpected error channel closing")
+					return nil, false, false, fmt.Errorf("unexpected error channel closing")
 				}
-				return nil, false, err
+				return nil, false, false, err
 			}
 		}
 	} else {
 		for index, word := range katakanaAndWordBits.Words {
 			newIndices := generateNewUsedIndices(usedIndices, index)
-			wordList, ok, err := i.f(word, newIndices, katakanaBitsAndWords, remainKatakanaBits)
+			wordList, ok, cacheUsed, err := i.f(word, newIndices, katakanaBitsAndWords, remainKatakanaBits)
 			if err != nil {
-				return nil, false, err
+				return nil, false, false, err
 			}
 			if ok {
 				irohaWordLists = append(irohaWordLists, wordList...)
 			}
-			i.log.PrintProgressLog(depth, "")
+			msg := ""
+			if cacheUsed {
+				msg = "cache used"
+			} else if depth <= i.depths.MaxStorage {
+				msg = "cache saved"
+			}
+			i.log.PrintProgressLog(depth, msg)
 		}
 	}
 
 	// どれも入れない場合
 	if remainKatakanaBits.has(katakanaAndWordBits.KatakanaBits) {
-		otherIrohaWordBitsLists, ok, err := i.searchByBits(append(usedIndices, -1), katakanaBitsAndWords[1:], remainKatakanaBits)
+		otherIrohaWordBitsLists, ok, cacheUsed, err := i.searchByBits(append(usedIndices, -1), katakanaBitsAndWords[1:], remainKatakanaBits)
 		if err != nil {
-			return nil, false, err
+			return nil, false, false, err
 		}
 		if ok {
 			irohaWordLists = append(irohaWordLists, otherIrohaWordBitsLists...)
 		}
-	}
+		msg := "no-add"
+		if cacheUsed {
+			msg += " / cache used"
+		} else if depth <= i.depths.MaxStorage {
+			msg += " / cache saved"
+		}
 
-	if err := i.storage.Set(usedIndices); err != nil {
-		return nil, false, err
+		i.log.PrintProgressLog(depth, msg)
 	}
 
-	return irohaWordLists, len(irohaWordLists) > 0, nil
+	if depth <= i.depths.MaxStorage {
+		if err := i.storage.Set(usedIndices, irohaWordLists); err != nil {
+			return nil, false, false, err
+		}
+	}
+	return irohaWordLists, len(irohaWordLists) > 0, false, nil
 }
 
 func generateNewUsedIndices(usedIndices []int, newIndex int) []int {

lib/storage.go

@@ -1,8 +1,18 @@
 package lib
 
+import (
+	"encoding/json"
+	"fmt"
+	"strconv"
+
+	"github.com/pkg/errors"
+
+	bolt "github.com/mpppk/bbolt"
+)
+
 type Storage interface {
-	Set(curs []int) error
-	Get(curs []int) ([][]*Word, bool, error)
+	Set(indices []int, wordsList [][]*Word) error
+	Get(indices []int) ([][]*Word, bool, error)
 }
 
 type NoStorage struct{}
@@ -11,6 +21,82 @@ func (e *NoStorage) Get(indices []int) ([][]*Word, bool, error) {
 	return nil, false, nil
 }
 
-func (e *NoStorage) Set(indices []int) error {
+func (e *NoStorage) Set(indices []int, wordsList [][]*Word) error {
 	return nil
 }
+
+type BoltStorage struct {
+	db         *bolt.DB
+	bucketName string
+}
+
+func NewBoltStorage(dbPath string) (*BoltStorage, error) {
+	db, err := bolt.Open(dbPath, 0600, nil)
+	if err != nil {
+		return nil, err
+	}
+	boltStorage := &BoltStorage{
+		db:         db,
+		bucketName: "main",
+	}
+	err = boltStorage.createBucketIfNotExists(boltStorage.bucketName)
+	return boltStorage, err
+}
+
+func (b *BoltStorage) Get(indices []int) (wordsList [][]*Word, ok bool, err error) {
+	wordsList = make([][]*Word, 0, 10)
+	err = b.db.View(func(tx *bolt.Tx) error {
+		bucket := tx.Bucket([]byte(b.bucketName))
+		if bucket == nil {
+			return fmt.Errorf("failed to retrive bucket(%s)", b.bucketName)
+		}
+		v := bucket.Get(toStorageKey(indices))
+		if v == nil {
+			ok = false
+			return nil
+		}
+		ok = true
+		return json.Unmarshal(v, &wordsList)
+	})
+	return
+}
+
+func (b *BoltStorage) Set(indices []int, wordsList [][]*Word) error {
+	wl := wordsList
+	if wl == nil {
+		wl = make([][]*Word, 0)
+	}
+
+	wordsListJsonBytes, err := json.Marshal(wl)
+
+	if err != nil {
+		return err
+	}
+	return b.db.Update(func(tx *bolt.Tx) error {
+		b := tx.Bucket([]byte(b.bucketName))
+		err := b.Put(
+			toStorageKey(indices),
+			wordsListJsonBytes)
+		return errors.Wrapf(err, "failed to put wordsList to bolt DB: indices:%s", indices)
+	})
+}
+
+func toStorageKey(indices []int) []byte {
+	strKey := ""
+	if len(indices) == 0 {
+		return []byte("no-index")
+	}
+	for _, index := range indices {
+		strKey += strconv.Itoa(index) + ":"
+	}
+	return []byte(strKey)
+}
+
+func (b *BoltStorage) createBucketIfNotExists(bucketName string) error {
+	return b.db.Update(func(tx *bolt.Tx) error {
+		if _, err := tx.CreateBucketIfNotExists([]byte(bucketName)); err != nil {
+			return fmt.Errorf("failed to create %s bucket: %s", bucketName, err)
+		}
+		return nil
+	})
+}