github.com/storacha/go-ucanto@v0.7.2/core/dag/blockstore/blockstore.go (about)

     1  package blockstore
     2  
     3  import (
     4  	"fmt"
     5  	"iter"
     6  	"sync"
     7  
     8  	"github.com/storacha/go-ucanto/core/ipld"
     9  )
    10  
    11  type BlockReader interface {
    12  	Get(link ipld.Link) (ipld.Block, bool, error)
    13  	Iterator() iter.Seq2[ipld.Block, error]
    14  }
    15  
    16  type BlockWriter interface {
    17  	Put(block ipld.Block) error
    18  }
    19  
    20  type BlockStore interface {
    21  	BlockReader
    22  	BlockWriter
    23  }
    24  
    25  type blockreader struct {
    26  	keys []string
    27  	blks map[string]ipld.Block
    28  }
    29  
    30  func (br *blockreader) Get(link ipld.Link) (ipld.Block, bool, error) {
    31  	b, ok := br.blks[link.String()]
    32  	return b, ok, nil
    33  }
    34  
    35  func (br *blockreader) Iterator() iter.Seq2[ipld.Block, error] {
    36  	return func(yield func(ipld.Block, error) bool) {
    37  		for _, k := range br.keys {
    38  			v, ok := br.blks[k]
    39  			var err error
    40  			if !ok {
    41  				err = fmt.Errorf("missing block for key: %s", k)
    42  			}
    43  			if !yield(v, err) {
    44  				return
    45  			}
    46  		}
    47  	}
    48  }
    49  
    50  type blockstore struct {
    51  	sync.RWMutex
    52  	blockreader
    53  }
    54  
    55  func (bs *blockstore) Put(block ipld.Block) error {
    56  	bs.RLock()
    57  	defer bs.RUnlock()
    58  
    59  	_, ok := bs.blks[block.Link().String()]
    60  	if ok {
    61  		return nil
    62  	}
    63  
    64  	bs.blks[block.Link().String()] = block
    65  	bs.keys = append(bs.keys, block.Link().String())
    66  
    67  	return nil
    68  }
    69  
    70  func (bs *blockstore) Get(link ipld.Link) (ipld.Block, bool, error) {
    71  	bs.Lock()
    72  	defer bs.Unlock()
    73  	return bs.blockreader.Get(link)
    74  }
    75  
    76  func (bs *blockstore) Iterator() iter.Seq2[ipld.Block, error] {
    77  	bs.Lock()
    78  	defer bs.Unlock()
    79  	return func(yield func(ipld.Block, error) bool) {
    80  		for _, k := range bs.keys {
    81  			v, ok := bs.blks[k]
    82  			var err error
    83  			if !ok {
    84  				err = fmt.Errorf("missing block for key: %s", k)
    85  			}
    86  			if !yield(v, err) {
    87  				return
    88  			}
    89  		}
    90  	}
    91  }
    92  
    93  // Option is an option configuring a block reader/writer.
    94  type Option func(cfg *bsConfig) error
    95  
    96  type bsConfig struct {
    97  	blks     []ipld.Block
    98  	blksiter iter.Seq2[ipld.Block, error]
    99  }
   100  
   101  // WithBlocks configures the blocks the blockstore should contain.
   102  func WithBlocks(blks []ipld.Block) Option {
   103  	return func(cfg *bsConfig) error {
   104  		cfg.blks = blks
   105  		return nil
   106  	}
   107  }
   108  
   109  // WithBlocksIterator configures the blocks the blockstore should contain.
   110  func WithBlocksIterator(blks iter.Seq2[ipld.Block, error]) Option {
   111  	return func(cfg *bsConfig) error {
   112  		cfg.blksiter = blks
   113  		return nil
   114  	}
   115  }
   116  
   117  func NewBlockStore(options ...Option) (BlockStore, error) {
   118  	cfg := bsConfig{}
   119  	for _, opt := range options {
   120  		if err := opt(&cfg); err != nil {
   121  			return nil, err
   122  		}
   123  	}
   124  	bs := &blockstore{
   125  		blockreader: blockreader{
   126  			keys: []string{},
   127  			blks: map[string]ipld.Block{},
   128  		},
   129  	}
   130  	for _, b := range cfg.blks {
   131  		err := bs.Put(b)
   132  		if err != nil {
   133  			return nil, err
   134  		}
   135  	}
   136  	if cfg.blksiter != nil {
   137  		for b, err := range cfg.blksiter {
   138  			if err != nil {
   139  				return nil, err
   140  			}
   141  			err := bs.Put(b)
   142  			if err != nil {
   143  				return nil, err
   144  			}
   145  		}
   146  	}
   147  	return bs, nil
   148  }
   149  
   150  func NewBlockReader(options ...Option) (BlockReader, error) {
   151  	cfg := bsConfig{}
   152  	for _, opt := range options {
   153  		if err := opt(&cfg); err != nil {
   154  			return nil, err
   155  		}
   156  	}
   157  
   158  	keys := []string{}
   159  	blks := map[string]ipld.Block{}
   160  
   161  	for _, b := range cfg.blks {
   162  		_, ok := blks[b.Link().String()]
   163  		if ok {
   164  			continue
   165  		}
   166  		blks[b.Link().String()] = b
   167  		keys = append(keys, b.Link().String())
   168  	}
   169  	if cfg.blksiter != nil {
   170  		for b, err := range cfg.blksiter {
   171  			if err != nil {
   172  				return nil, err
   173  			}
   174  			_, ok := blks[b.Link().String()]
   175  			if ok {
   176  				continue
   177  			}
   178  			blks[b.Link().String()] = b
   179  			keys = append(keys, b.Link().String())
   180  		}
   181  	}
   182  
   183  	return &blockreader{keys, blks}, nil
   184  }
   185  
   186  func WriteInto(view ipld.View, bs BlockWriter) error {
   187  	for b := range view.Blocks() {
   188  		err := bs.Put(b)
   189  		if err != nil {
   190  			return fmt.Errorf("putting proof block: %w", err)
   191  		}
   192  	}
   193  	return nil
   194  }