github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/mapio/map.go (about)

     1  // Copyright 2018 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package mapio
     6  
     7  import (
     8  	"errors"
     9  	"io"
    10  	"sync"
    11  )
    12  
    13  // Map is a read-only, sorted map backed by an io.ReadSeeker. The
    14  // on-disk layout of maps are described by the package documentation.
    15  // Maps support both lookup and (ordered) iteration. A Map instance
    16  // maintains a current position, starting out at the first entry.
    17  type Map struct {
    18  	mu    sync.Mutex
    19  	r     io.ReadSeeker
    20  	index block
    21  }
    22  
    23  // New opens the map at the provided io.ReadSeeker (usually a file).
    24  func New(r io.ReadSeeker) (*Map, error) {
    25  	m := &Map{r: r}
    26  	return m, m.init()
    27  }
    28  
    29  func (m *Map) init() error {
    30  	if _, err := m.r.Seek(-mapTrailerSize, io.SeekEnd); err != nil {
    31  		return err
    32  	}
    33  	trailer := make([]byte, mapTrailerSize)
    34  	if _, err := io.ReadFull(m.r, trailer); err != nil {
    35  		return err
    36  	}
    37  	metaAddr, _ := getBlockAddr(trailer)
    38  	if metaAddr != (blockAddr{}) {
    39  		return errors.New("non-empty meta block index")
    40  	}
    41  	indexAddr, _ := getBlockAddr(trailer[maxBlockAddrSize:])
    42  	magic := order.Uint64(trailer[len(trailer)-8:])
    43  	if magic != mapTrailerMagic {
    44  		return errors.New("wrong magic")
    45  	}
    46  	if err := m.readBlock(indexAddr, &m.index); err != nil {
    47  		return err
    48  	}
    49  	if !m.index.Scan() {
    50  		return errors.New("empty index")
    51  	}
    52  	return nil
    53  }
    54  
    55  func (m *Map) readBlock(addr blockAddr, block *block) error {
    56  	if block.p != nil && cap(block.p) >= int(addr.len) {
    57  		block.p = block.p[:addr.len]
    58  	} else {
    59  		block.p = make([]byte, addr.len)
    60  	}
    61  	m.mu.Lock()
    62  	defer m.mu.Unlock()
    63  	if _, err := m.r.Seek(int64(addr.off), io.SeekStart); err != nil {
    64  		return err
    65  	}
    66  	if _, err := io.ReadFull(m.r, block.p); err != nil {
    67  		return err
    68  	}
    69  	return block.init()
    70  }
    71  
    72  // Seek returns a map scanner beginning at the first key in the map
    73  // >= the provided key.
    74  func (m *Map) Seek(key []byte) *MapScanner {
    75  	s := &MapScanner{parent: m, index: m.index}
    76  	s.index.Seek(key)
    77  	if s.index.Scan() {
    78  		addr, _ := getBlockAddr(s.index.Value())
    79  		if s.err = m.readBlock(addr, &s.data); s.err == nil {
    80  			s.data.Seek(key)
    81  		}
    82  	}
    83  	return s
    84  }
    85  
    86  // MapScanner implements ordered iteration over a map.
    87  type MapScanner struct {
    88  	parent      *Map
    89  	err         error
    90  	data, index block
    91  }
    92  
    93  // Scan scans the next entry, returning true on success. When Scan
    94  // returns false, the caller should inspect Err to distinguish
    95  // between scan completion and scan error.
    96  func (m *MapScanner) Scan() bool {
    97  	for m.err == nil && !m.data.Scan() {
    98  		if !m.index.Scan() {
    99  			return false
   100  		}
   101  		addr, _ := getBlockAddr(m.index.Value())
   102  		m.err = m.parent.readBlock(addr, &m.data)
   103  	}
   104  	return m.err == nil
   105  }
   106  
   107  // Err returns the last error encountered while scanning.
   108  func (m *MapScanner) Err() error {
   109  	return m.err
   110  }
   111  
   112  // Key returns the key that was last scanned.
   113  func (m *MapScanner) Key() []byte {
   114  	return m.data.Key()
   115  }
   116  
   117  // Value returns the value that was last scanned.
   118  func (m *MapScanner) Value() []byte {
   119  	return m.data.Value()
   120  }