Prism VCS > Tree [f4699a2]

/object_store.go/

// Copyright (c) 2025, Christian Lee Seibold
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
//    contributors may be used to endorse or promote products derived from
//    this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package main

import (
	"errors"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/klauspost/compress/zstd"
)

type ObjectType uint

const (
	ObjectType_Unknown ObjectType = iota
	ObjectType_Commit
	ObjectType_Tree
	ObjectType_Blob
)

type ObjectStore struct {
	vcsDirectory string
}

func (s *ObjectStore) GetBranches() ([]string, error) {
	entries, err := os.ReadDir(filepath.Join(s.vcsDirectory + "/refs/branches/"))
	if err != nil {
		return nil, err
	}

	var branchNames []string
	for _, entry := range entries {
		branchNames = append(branchNames, entry.Name())
	}
	return branchNames, nil
}

// Gets an ObjectReader to read the object uncompressed. Object must be closed after using.
func (s *ObjectStore) GetObject(hash Hash) (*ObjectReader, error) {
	path := filepath.Join(s.vcsDirectory, "objects", hash.HexString())
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}

	decoder, err := zstd.NewReader(file)
	if err != nil {
		return nil, err
	}

	// Read header
	header := strings.Builder{}
	var b = [1]byte{}
	for {
		_, err = decoder.Read(b[:])
		if err != nil && err != io.EOF {
			break
		}

		if b[0] == byte('\000') {
			break
		}
		header.WriteByte(b[0])
	}

	typeString, sizeString, _ := strings.Cut(header.String(), " ")
	objType := ObjectType_Unknown
	switch typeString {
	case "blob":
		objType = ObjectType_Blob
	case "tree":
		objType = ObjectType_Tree
	case "commit":
		objType = ObjectType_Commit
	}

	size, _ := strconv.ParseUint(sizeString, 10, 64)
	return &ObjectReader{Type: objType, Size: size, f: file, Decoder: decoder}, nil
}

// Generic object to be read from
type ObjectReader struct {
	Type ObjectType
	Size uint64
	f    *os.File
	*zstd.Decoder
}

// Closes the compression decoder and the file.
func (o ObjectReader) Close() error {
	o.IOReadCloser().Close()
	return o.f.Close()
}

var ErrObjectExists = errors.New("Object already exists.")

// Write an object to the objects store from the reader, automatically handling the compression.
// Returns ErrObjectExists if an object with the given hash already exists.
func (s *ObjectStore) WriteObject(hash Hash, reader io.Reader) error {
	path := filepath.Join(s.vcsDirectory, "objects", hash.HexString())

	// Return error if object already exists.
	if _, err := os.Stat(path); err == nil {
		return ErrObjectExists
	}

	file, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR|os.O_EXCL, 0644)
	if err != nil {
		return err
	}
	zw, err := zstd.NewWriter(file, zstd.WithEncoderLevel(zstd.SpeedBestCompression))
	if err != nil {
		return err
	}
	_, err = io.Copy(zw, reader)
	if err != nil {
		return err
	}
	file.Sync()

	zw.Close()
	file.Close()
	return nil
}