goncdu/node/saved.go

package node

import (
	"compress/gzip"
	"encoding/json"
	"errors"
	"io"
	"os"
)

type StoreType int

const (
	JSON_GZ StoreType = iota
	JSON_NO_IDENT
	JSON_IDENT
)

// MarshalJSON 实现自定义序列化
func (pn *PathNode) MarshalJSON() ([]byte, error) {
	type Alias PathNode
	return json.Marshal(&struct {
		*Alias
		Path      string               `json:"path"`
		Children  map[string]*PathNode `json:"children"`
		Count     int                  `json:"count"`
		Node_type PathNodeType         `json:"type"`
		Size      int64                `json:"size"`
	}{
		Alias:     (*Alias)(pn),
		Path:      pn.path,
		Children:  pn.children,
		Count:     pn.count,
		Size:      pn.size,
		Node_type: pn.node_type,
	})
}

// UnmarshalJSON 实现自定义反序列化
func (pn *PathNode) UnmarshalJSON(data []byte) error {
	type Alias PathNode
	temp := &struct {
		*Alias
		Path      string               `json:"path"`
		Children  map[string]*PathNode `json:"children"`
		Count     int                  `json:"count"`
		Node_type PathNodeType         `json:"type"`
		Size      int64                `json:"size"`
	}{
		Alias: (*Alias)(pn),
	}

	if err := json.Unmarshal(data, &temp); err != nil {
		return err
	}

	pn.path = temp.Path
	pn.children = temp.Children
	pn.count = temp.Count
	pn.size = temp.Size
	pn.node_type = temp.Node_type
	return nil
}

// ToJSON 将 PathNode 序列化为 JSON 字符串
func (pn *PathNode) ToJSON(savedPath string, savedType StoreType) error {
	f, err := os.Create(savedPath)
	if err != nil {
		return err
	}
	defer f.Close()

	switch savedType {
	case JSON_GZ:
		data, err := json.Marshal(pn)
		if err != nil {
			return err
		}
		gz := gzip.NewWriter(f)
		defer gz.Close()
		if _, err := gz.Write(data); err != nil {
			return err
		}
	case JSON_NO_IDENT:
		data, err := json.Marshal(pn)
		if err != nil {
			return err
		}
		if _, err = f.Write(data); err != nil {
			return err
		}
	case JSON_IDENT:
		data, err := json.MarshalIndent(pn, "", "  ")
		if err != nil {
			return err
		}
		if _, err = f.Write(data); err != nil {
			return err
		}
	}
	return nil
}

// FromJSON 从 JSON 字符串反序列化为 PathNode
func (pn *PathNode) FromJSON(savedPath string, savedType StoreType) error {
	if pn.parent != nil {
		return errors.New("PathNode.Parent should be nil")
	}

	f, err := os.Open(savedPath)
	if err != nil {
		return err
	}
	defer f.Close()

	var jsonData []byte
	switch savedType {
	case JSON_GZ:
		gzr, err := gzip.NewReader(f)
		if err != nil {
			return err
		}
		defer gzr.Close()
		if jsonData, err = io.ReadAll(gzr); err != nil {
			return err
		}
	default:
		// case JSON_IDENT:
		// case JSON_NO_IDENT:
		if jsonData, err = io.ReadAll(f); err != nil {
			return err
		}
	}

	pn.ClearWithoutFlush()
	if err := json.Unmarshal([]byte(jsonData), pn); err != nil {
		return err
	}

	// RestoreParentLinks 恢复 Parent 关系
	var restoreParentLinks func(pn *PathNode)
	restoreParentLinks = func(pn *PathNode) {
		for _, child := range pn.children {
			child.parent = pn
			restoreParentLinks(child)
		}
	}

	restoreParentLinks(pn)
	return nil
}