internal/pkg: move IR primitives out of cache

These are memory management and caching primitives. Having them as part of Cache is cumbersome and requires a temporary directory that is never used. This change isolates them from Cache to enable independent use.

Signed-off-by: Ophestra <cat@gensokyo.uk>

2 files changed, 62 insertions, 41 deletions

diff --git a/internal/pkg/ir.go b/internal/pkg/ir.go
index 21ba8edc..41e75845 100644
--- a/internal/pkg/ir.go
+++ b/internal/pkg/ir.go
@@ -3,7 +3,6 @@ package pkg
 import (
 	"bufio"
 	"bytes"
-	"context"
 	"crypto/sha512"
 	"encoding/binary"
 	"errors"
@@ -11,6 +10,7 @@ import (
 	"io"
 	"slices"
 	"strconv"
+	"sync"
 	"syscall"
 	"unique"
 	"unsafe"
@@ -39,22 +39,45 @@ func panicToError(errP *error) {
 	}
 }
 
+// irCache implements [IRCache].
+type irCache struct {
+	// Artifact to [unique.Handle] of identifier cache.
+	artifact sync.Map
+	// Identifier free list, must not be accessed directly.
+	identPool sync.Pool
+}
+
+// zeroIRCache returns the initialised value of irCache.
+func zeroIRCache() irCache {
+	return irCache{
+		identPool: sync.Pool{New: func() any { return new(extIdent) }},
+	}
+}
+
+// IRCache provides memory management and caching primitives for IR and
+// identifier operations against [Artifact] implementations.
+//
+// The zero value is not safe for use.
+type IRCache struct{ irCache }
+
+// NewIR returns the address of a new [IRCache].
+func NewIR() *IRCache {
+	return &IRCache{zeroIRCache()}
+}
+
 // IContext is passed to [Artifact.Params] and provides methods for writing
 // values to the IR writer. It does not expose the underlying [io.Writer].
 //
 // IContext is valid until [Artifact.Params] returns.
 type IContext struct {
-	// Address of underlying [Cache], should be zeroed or made unusable after
+	// Address of underlying irCache, should be zeroed or made unusable after
 	// [Artifact.Params] returns and must not be exposed directly.
-	cache *Cache
+	ic *irCache
 	// Written to by various methods, should be zeroed after [Artifact.Params]
 	// returns and must not be exposed directly.
 	w io.Writer
 }
 
-// Unwrap returns the underlying [context.Context].
-func (i *IContext) Unwrap() context.Context { return i.cache.ctx }
-
 // irZero is a zero IR word.
 var irZero [wordSize]byte
 
@@ -136,11 +159,11 @@ func (i *IContext) mustWrite(p []byte) {
 // WriteIdent is not defined for an [Artifact] not part of the slice returned by
 // [Artifact.Dependencies].
 func (i *IContext) WriteIdent(a Artifact) {
-	buf := i.cache.getIdentBuf()
-	defer i.cache.putIdentBuf(buf)
+	buf := i.ic.getIdentBuf()
+	defer i.ic.putIdentBuf(buf)
 
 	IRKindIdent.encodeHeader(0).put(buf[:])
-	*(*ID)(buf[wordSize:]) = i.cache.Ident(a).Value()
+	*(*ID)(buf[wordSize:]) = i.ic.Ident(a).Value()
 	i.mustWrite(buf[:])
 }
 
@@ -183,19 +206,19 @@ func (i *IContext) WriteString(s string) {
 
 // Encode writes a deterministic, efficient representation of a to w and returns
 // the first non-nil error encountered while writing to w.
-func (c *Cache) Encode(w io.Writer, a Artifact) (err error) {
+func (ic *irCache) Encode(w io.Writer, a Artifact) (err error) {
 	deps := a.Dependencies()
 	idents := make([]*extIdent, len(deps))
 	for i, d := range deps {
-		dbuf, did := c.unsafeIdent(d, true)
+		dbuf, did := ic.unsafeIdent(d, true)
 		if dbuf == nil {
-			dbuf = c.getIdentBuf()
+			dbuf = ic.getIdentBuf()
 			binary.LittleEndian.PutUint64(dbuf[:], uint64(d.Kind()))
 			*(*ID)(dbuf[wordSize:]) = did.Value()
 		} else {
-			c.storeIdent(d, dbuf)
+			ic.storeIdent(d, dbuf)
 		}
-		defer c.putIdentBuf(dbuf)
+		defer ic.putIdentBuf(dbuf)
 		idents[i] = dbuf
 	}
 	slices.SortFunc(idents, func(a, b *extIdent) int {
@@ -221,10 +244,10 @@ func (c *Cache) Encode(w io.Writer, a Artifact) (err error) {
 	}
 
 	func() {
-		i := IContext{c, w}
+		i := IContext{ic, w}
 
 		defer panicToError(&err)
-		defer func() { i.cache, i.w = nil, nil }()
+		defer func() { i.ic, i.w = nil, nil }()
 
 		a.Params(&i)
 	}()
@@ -233,7 +256,7 @@ func (c *Cache) Encode(w io.Writer, a Artifact) (err error) {
 	}
 
 	var f IREndFlag
-	kcBuf := c.getIdentBuf()
+	kcBuf := ic.getIdentBuf()
 	sz := wordSize
 	if kc, ok := a.(KnownChecksum); ok {
 		f |= IREndKnownChecksum
@@ -243,13 +266,13 @@ func (c *Cache) Encode(w io.Writer, a Artifact) (err error) {
 	IRKindEnd.encodeHeader(uint32(f)).put(kcBuf[:])
 
 	_, err = w.Write(kcBuf[:sz])
-	c.putIdentBuf(kcBuf)
+	ic.putIdentBuf(kcBuf)
 	return
 }
 
 // encodeAll implements EncodeAll by recursively encoding dependencies and
 // performs deduplication by value via the encoded map.
-func (c *Cache) encodeAll(
+func (ic *irCache) encodeAll(
 	w io.Writer,
 	a Artifact,
 	encoded map[Artifact]struct{},
@@ -259,13 +282,13 @@ func (c *Cache) encodeAll(
 	}
 
 	for _, d := range a.Dependencies() {
-		if err = c.encodeAll(w, d, encoded); err != nil {
+		if err = ic.encodeAll(w, d, encoded); err != nil {
 			return
 		}
 	}
 
 	encoded[a] = struct{}{}
-	return c.Encode(w, a)
+	return ic.Encode(w, a)
 }
 
 // EncodeAll writes a self-describing IR stream of a to w and returns the first
@@ -283,8 +306,8 @@ func (c *Cache) encodeAll(
 // the ident cache, nor does it contribute identifiers it computes back to the
 // ident cache. Because of this, multiple invocations of EncodeAll will have
 // similar cost and does not amortise when combined with a call to Cure.
-func (c *Cache) EncodeAll(w io.Writer, a Artifact) error {
-	return c.encodeAll(w, a, make(map[Artifact]struct{}))
+func (ic *irCache) EncodeAll(w io.Writer, a Artifact) error {
+	return ic.encodeAll(w, a, make(map[Artifact]struct{}))
 }
 
 // ErrRemainingIR is returned for a [IRReadFunc] that failed to call
diff --git a/internal/pkg/pkg.go b/internal/pkg/pkg.go
index b2eceb25..fea3147f 100644
--- a/internal/pkg/pkg.go
+++ b/internal/pkg/pkg.go
@@ -559,10 +559,8 @@ type Cache struct {
 	// Immutable job count, when applicable.
 	jobs int
 
-	// Artifact to [unique.Handle] of identifier cache.
-	artifact sync.Map
-	// Identifier free list, must not be accessed directly.
-	identPool sync.Pool
+	// Must not be exposed directly.
+	irCache
 
 	// Synchronises access to dirChecksum.
 	checksumMu sync.RWMutex
@@ -594,24 +592,24 @@ type Cache struct {
 type extIdent [wordSize + len(ID{})]byte
 
 // getIdentBuf returns the address of an extIdent for Ident.
-func (c *Cache) getIdentBuf() *extIdent { return c.identPool.Get().(*extIdent) }
+func (ic *irCache) getIdentBuf() *extIdent { return ic.identPool.Get().(*extIdent) }
 
 // putIdentBuf adds buf to identPool.
-func (c *Cache) putIdentBuf(buf *extIdent) { c.identPool.Put(buf) }
+func (ic *irCache) putIdentBuf(buf *extIdent) { ic.identPool.Put(buf) }
 
 // storeIdent adds an [Artifact] to the artifact cache.
-func (c *Cache) storeIdent(a Artifact, buf *extIdent) unique.Handle[ID] {
+func (ic *irCache) storeIdent(a Artifact, buf *extIdent) unique.Handle[ID] {
 	idu := unique.Make(ID(buf[wordSize:]))
-	c.artifact.Store(a, idu)
+	ic.artifact.Store(a, idu)
 	return idu
 }
 
 // Ident returns the identifier of an [Artifact].
-func (c *Cache) Ident(a Artifact) unique.Handle[ID] {
-	buf, idu := c.unsafeIdent(a, false)
+func (ic *irCache) Ident(a Artifact) unique.Handle[ID] {
+	buf, idu := ic.unsafeIdent(a, false)
 	if buf != nil {
-		idu = c.storeIdent(a, buf)
-		c.putIdentBuf(buf)
+		idu = ic.storeIdent(a, buf)
+		ic.putIdentBuf(buf)
 	}
 	return idu
 }
@@ -619,17 +617,17 @@ func (c *Cache) Ident(a Artifact) unique.Handle[ID] {
 // unsafeIdent implements Ident but returns the underlying buffer for a newly
 // computed identifier. Callers must return this buffer to identPool. encodeKind
 // is only a hint, kind may still be encoded in the buffer.
-func (c *Cache) unsafeIdent(a Artifact, encodeKind bool) (
+func (ic *irCache) unsafeIdent(a Artifact, encodeKind bool) (
 	buf *extIdent,
 	idu unique.Handle[ID],
 ) {
-	if id, ok := c.artifact.Load(a); ok {
+	if id, ok := ic.artifact.Load(a); ok {
 		idu = id.(unique.Handle[ID])
 		return
 	}
 
 	if ki, ok := a.(KnownIdent); ok {
-		buf = c.getIdentBuf()
+		buf = ic.getIdentBuf()
 		if encodeKind {
 			binary.LittleEndian.PutUint64(buf[:], uint64(a.Kind()))
 		}
@@ -637,9 +635,9 @@ func (c *Cache) unsafeIdent(a Artifact, encodeKind bool) (
 		return
 	}
 
-	buf = c.getIdentBuf()
+	buf = ic.getIdentBuf()
 	h := sha512.New384()
-	if err := c.Encode(h, a); err != nil {
+	if err := ic.Encode(h, a); err != nil {
 		// unreachable
 		panic(err)
 	}
@@ -1876,7 +1874,7 @@ func open(
 		msg:  msg,
 		base: base,
 
-		identPool: sync.Pool{New: func() any { return new(extIdent) }},
+		irCache: zeroIRCache(),
 
 		ident:        make(map[unique.Handle[ID]]unique.Handle[Checksum]),
 		identErr:     make(map[unique.Handle[ID]]error),