prime-div/io.go

94 lines
1.7 KiB
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)
}
}