1 files changed, 41 insertions, 10 deletions
diff --git a/core/index/index.go b/core/index/index.go
index 2e669fc..b00f524 100644
--- a/core/index/index.go
+++ b/core/index/index.go
@@ -1,6 +1,7 @@
 package index
 
 import (
+	"bufio"
 	"encoding/gob"
 	"os"
 	"searchEngine/core/analyzer"
@@ -14,9 +15,10 @@ type Posting struct {
 }
 
 type InvertedIndex struct {
-	Mu       sync.RWMutex `gob:"-"`
-	Data     map[string][]Posting
-	DocNames map[int]string
+	Mu         sync.RWMutex `gob:"-"`
+	Data       map[string][]Posting
+	DocNames   map[int]string
+	DocLengths map[int]int
 }
 
 func New() *InvertedIndex {
@@ -51,7 +53,7 @@ func (idx *InvertedIndex) Add(docID int, title string, bodyTokens []string) {
 	defer idx.Mu.Unlock()
 
 	idx.DocNames[docID] = title
-
+	idx.DocLengths[docID] = len(bodyTokens)
 	for token, stats := range docStats {
 		idx.Data[token] = append(idx.Data[token], Posting{
 			DocId:          docID,
@@ -71,8 +73,25 @@ func (idx *InvertedIndex) Save(filepath string) error {
 	}
 	defer file.Close()
 
-	encoder := gob.NewEncoder(file)
-	return encoder.Encode(idx)
+	writer := bufio.NewWriter(file)
+
+	dto := struct {
+		Data       map[string][]Posting
+		DocNames   map[int]string
+		DocLengths map[int]int
+	}{
+		Data:       idx.Data,
+		DocNames:   idx.DocNames,
+		DocLengths: idx.DocLengths,
+	}
+
+	encoder := gob.NewEncoder(writer)
+	err = encoder.Encode(dto)
+
+	if flushErr := writer.Flush(); flushErr != nil {
+		return flushErr
+	}
+	return err
 }
 
 func Load(filepath string) (*InvertedIndex, error) {
@@ -82,11 +101,23 @@ func Load(filepath string) (*InvertedIndex, error) {
 	}
 	defer file.Close()
 
-	idx := &InvertedIndex{}
-	decoder := gob.NewDecoder(file)
-	err = decoder.Decode(idx)
+	reader := bufio.NewReader(file)
+
+	var dto struct {
+		Data       map[string][]Posting
+		DocNames   map[int]string
+		DocLengths map[int]int
+	}
+
+	decoder := gob.NewDecoder(reader)
+	err = decoder.Decode(&dto)
 	if err != nil {
 		return nil, err
 	}
-	return idx, nil
+
+	return &InvertedIndex{
+		Data:       dto.Data,
+		DocNames:   dto.DocNames,
+		DocLengths: dto.DocLengths,
+	}, nil
 }