github.com/coyove/common@v0.0.0-20240403014525-f70e643f9de8/quadtree/mgr.go (about)

     1  package quadtree
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"strconv"
     7  	"sync"
     8  )
     9  
    10  type Database interface {
    11  	Load(id string) (QuadTree, error)
    12  	Store(t QuadTree) error
    13  	StoreElement(id string, e Element) error
    14  	DeleteAllElements(id string) error
    15  	DeleteElement(id string, e Element) error
    16  	StoreOrthant(id string, o int, oid string) (existed bool, err error)
    17  }
    18  
    19  type MemoryDatabase struct {
    20  	sync.Mutex
    21  	m map[string]map[string][]byte
    22  	e map[string]map[Point]Element
    23  }
    24  
    25  func NewMemoryDatabase() *MemoryDatabase {
    26  	return &MemoryDatabase{
    27  		m: map[string]map[string][]byte{},
    28  		e: map[string]map[Point]Element{},
    29  	}
    30  }
    31  
    32  func (m *MemoryDatabase) Load(id string) (QuadTree, error) {
    33  	m.Lock()
    34  	h := m.m[id]
    35  	m.Unlock()
    36  	if h == nil {
    37  		return QuadTree{}, fmt.Errorf("%q not found", id)
    38  	}
    39  	t := QuadTree{}
    40  	if err := json.Unmarshal(h["t"], &t); err != nil {
    41  		return QuadTree{}, err
    42  	}
    43  	m.Lock()
    44  	t.O[0], t.O[1], t.O[2], t.O[3] = string(h["0"]), string(h["1"]), string(h["2"]), string(h["3"])
    45  	t.Elems = m.e[id]
    46  	m.Unlock()
    47  	return t, nil
    48  }
    49  
    50  func (m *MemoryDatabase) Store(t QuadTree) error {
    51  	buf, _ := json.Marshal(t)
    52  	m.Lock()
    53  	m.m[t.ID] = map[string][]byte{"t": buf}
    54  	m.Unlock()
    55  	return nil
    56  }
    57  
    58  func (m *MemoryDatabase) StoreElement(id string, e Element) error {
    59  	m.Lock()
    60  	h := m.e[id]
    61  	if h == nil {
    62  		h = map[Point]Element{}
    63  		m.e[id] = h
    64  	}
    65  	h[e.Point] = e
    66  	m.Unlock()
    67  	return nil
    68  }
    69  
    70  func (m *MemoryDatabase) DeleteAllElements(id string) error {
    71  	m.Lock()
    72  	delete(m.e, id)
    73  	m.Unlock()
    74  	return nil
    75  }
    76  
    77  func (m *MemoryDatabase) DeleteElement(id string, e Element) error {
    78  	m.Lock()
    79  	h := m.e[id]
    80  	delete(h, e.Point)
    81  	m.Unlock()
    82  	return nil
    83  }
    84  
    85  func (m *MemoryDatabase) StoreOrthant(id string, o int, oid string) (existed bool, err error) {
    86  	m.Lock()
    87  	defer m.Unlock()
    88  	h := m.m[id]
    89  	if m == nil {
    90  		return false, fmt.Errorf("store %q orthant #%d: not found", id, o)
    91  	}
    92  	if _, exist := h[strconv.Itoa(o)]; exist {
    93  		return true, nil
    94  	}
    95  	h[strconv.Itoa(o)] = []byte(oid)
    96  	return false, nil
    97  }