prime-div/io.go

package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"strconv"
	"time"

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

func loadPrimes(useTxt bool) []int64 {
	if useTxt {
		return loadPrimesTxt()
	} else {
		return loadPrimesMsgpack()
	}
}

func loadPrimesTxt() []int64 {
	file := must(os.Open("prime.txt"))
	defer file.Close()

	var primes []int64

	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		nextPrime := must(strconv.ParseInt(scanner.Text(), 10, 64))
		primes = append(primes, nextPrime)
	}

	return primes
}

func loadPrimesMsgpack() []int64 {
	file := must(os.Open("prime.msgpack.zst"))
	defer file.Close()

	zstdDecoder := must(zstd.NewReader(file))
	defer zstdDecoder.Close()

	var primes []int64
	msgpackDecoder := msgpack.NewDecoder(zstdDecoder)
	err := msgpackDecoder.Decode(&primes)
	if err != nil {
		log.Fatal(err)
	}

	return primes
}

func writePrimes(primes []int64, useTxt bool) {
	startTime := time.Now()

	if useTxt {
		writePrimesTxt(primes)
	} else {
		writePrimesMsgpack(primes)
	}

	fmt.Printf("File writing took %v\n", time.Now().Sub(startTime))
}

func writePrimesTxt(primes []int64) {
	file := must(os.Create("prime.txt"))
	defer file.Close()

	writer := bufio.NewWriter(file)

	for _, prime := range primes {
		writer.WriteString(fmt.Sprintf("%d\n", prime))
	}

	writer.Flush()
}

func writePrimesMsgpack(primes []int64) {
	file := must(os.Create("prime.msgpack.zst"))
	defer file.Close()

	zstdEncoder := must(zstd.NewWriter(file, zstd.WithEncoderLevel(zstd.SpeedBestCompression)))
	defer zstdEncoder.Close()
	
	msgpackEncoder := msgpack.NewEncoder(zstdEncoder)
	err := msgpackEncoder.Encode(&primes)
	if err != nil {
		log.Fatal(err)
	}
}