notea-de-me/server/tree.go

package main

import (
	"io"
	"io/ioutil"
	"os"
	"path"

	yaml "gopkg.in/yaml.v2"
)

type Branch struct {
	Leaf     Leaf
	Branches map[ID]Branch
}

func (branch Branch) IsZero() bool {
	return branch.Leaf == (Leaf{}) && len(branch.Branches) == 0
}

func (branch Branch) ForEach(foo func(ID, Leaf) error) error {
	return branch.forEach(NewID(""), foo)
}

func (branch Branch) forEach(preid ID, foo func(ID, Leaf) error) error {
	if err := foo(preid, branch.Leaf); err != nil {
		return err
	}
	for id, child := range branch.Branches {
		if err := child.forEach(id, foo); err != nil {
			return err
		}
	}
	return nil
}

type Leaf struct {
	Title   string
	Deleted bool
	Content string
}

func (base Leaf) Merge(updated Leaf) Leaf {
	if updated.Title != "" {
		base.Title = updated.Title
	}
	if base.Title == "" {
		base.Title = "Untitled"
	}
	base.Deleted = updated.Deleted
	base.Content = updated.Content
	return base
}

type Tree struct {
	root       string
	cachedRoot Branch
}

func NewTree(root string) Tree {
	return Tree{root: root}
}

func (tree Tree) WithRoot(root string) Tree {
	tree.root = root
	tree.cachedRoot = Branch{}
	return tree
}

func (tree Tree) GetRootMeta() (Branch, error) {
	return tree.getRoot(NewID(""), false, false)
}

func (tree Tree) GetRoot() (Branch, error) {
	if !tree.cachedRoot.IsZero() {
		return tree.cachedRoot, nil
	}
	got, err := tree.getRoot(NewID(""), true, false)
	if err == nil {
		tree.cachedRoot = got
	}
	return got, err
}

func (tree Tree) getRoot(pid ID, withContent, withDeleted bool) (Branch, error) {
	m := Branch{Branches: map[ID]Branch{}}
	entries, err := os.ReadDir(tree.root)
	if os.IsNotExist(err) {
		return m, nil
	}
	if err != nil {
		return Branch{}, err
	}
	for _, entry := range entries {
		if entry.Name() == "data.yaml" {
			if b, err := peekFile(withContent, path.Join(tree.root, entry.Name())); err != nil {
				return Branch{}, err
			} else if err := yaml.Unmarshal(b, &m.Leaf); err != nil {
				return Branch{}, err
			}
			if !withContent {
				m.Leaf.Content = ""
			}
			if m.Leaf.Deleted && !withDeleted {
				return m, nil
			}
		} else if entry.IsDir() {
			subtree := tree.WithRoot(path.Join(tree.root, entry.Name()))
			if branch, err := subtree.getRoot(pid.Push(entry.Name()), withContent, withDeleted); err != nil {
				return Branch{}, err
			} else if !branch.IsZero() && (!branch.Leaf.Deleted || withDeleted) {
				m.Branches[pid.Push(entry.Name())] = branch
			}
		}
	}
	return m, nil
}

func peekFile(all bool, path string) ([]byte, error) {
	if !all {
		return ioutil.ReadFile(path)
	}
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	r := io.LimitReader(f, 1024)
	return ioutil.ReadAll(r)
}

func (tree Tree) toDir(id ID) string {
	return path.Dir(tree.toData(id))
}

func (tree Tree) toData(id ID) string {
	return path.Join(tree.root, string(id), "data.yaml")
}

func (tree Tree) Put(id ID, input Leaf) error {
	if _, err := os.Stat(tree.toData(id)); os.IsNotExist(err) {
		b, err := yaml.Marshal(Leaf{})
		if err != nil {
			return err
		}
		if err := ensureAndWrite(tree.toData(id), b); err != nil {
			return err
		}
	}
	old, err := tree.Get(id)
	if err != nil {
		return err
	}
	b, err := yaml.Marshal(old.Merge(input))
	if err != nil {
		return err
	}
	if err := ensureAndWrite(tree.toData(id), b); err != nil {
		return err
	}
	tree.cachedRoot = Branch{}
	return nil
}

func (tree Tree) Del(id ID) error {
	got, err := tree.Get(id)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return err
	}
	if got.Deleted {
		return nil
	}
	got.Deleted = true
	return tree.Put(id, got)
}

func (tree Tree) HardDel(id ID) error {
	os.RemoveAll(tree.toDir(id))
	tree.cachedRoot = Branch{}
	return nil
}

func (tree Tree) Get(id ID) (Leaf, error) {
	f, err := os.Open(tree.toData(id))
	if err != nil {
		return Leaf{}, err
	}
	defer f.Close()
	var got Leaf
	err = yaml.NewDecoder(f).Decode(&got)
	return got, err
}