inject.go

package inject

import (
	"errors"
	"fmt"
	"reflect"
	"slices"
	"sync"
)

type Injector interface {
	Get(reflect.Type) (reflect.Value, error)
}

type Injectable interface {
	Type() reflect.Type
	Inject(Injector) (reflect.Value, error)
}

func New(args ...any) (*Scope, error) {
	return (*Scope)(nil).Scope(args...)
}

func Get[T any](inj Injector) (t T, err error) {
	typ := reflect.TypeOf(t)
	var rt reflect.Value
	rt, err = inj.Get(typ)
	if err != nil {
		return
	}
	t = rt.Interface().(T)
	return
}

func Inject[T any](j Injector, fn interface{}) (t T, err error) {
	ttyp := reflect.TypeOf(&t).Elem()
	fnval := reflect.ValueOf(fn)
	fntyp := fnval.Type()
	if err = checkFunc(fnval, fntyp); err != nil {
		return
	}
	if fntyp.Out(0) != ttyp {
		err = fmt.Errorf("return type does not match generic arg: expected: %s, got: %s", ttyp.String(), fntyp.Out(0).String())
		return
	}
	var rt reflect.Value
	rt, err = inject(j, fnval)
	if err == nil {
		t = rt.Interface().(T)
	}
	return
}

func inject(j Injector, fn reflect.Value) (val reflect.Value, err error) {
	typ := fn.Type()

	if err = checkFunc(fn, typ); err != nil {
		return
	}

	// build the arguments and call the function
	args := make([]reflect.Value, typ.NumIn())
	for i := range args {
		var arg reflect.Value
		arg, err = j.Get(typ.In(i))
		if err != nil {
			err = errors.Join(fmt.Errorf("failed to get type '%s' from injector", typ.In(i).String()), err)
			return
		}
		args[i] = arg
	}
	results := fn.Call(args)

	// extract the optional error
	if len(results) == 2 {
		// we already know that results[1] has type error
		if i := results[1].Interface(); i != nil {
			err = fmt.Errorf("injected function returned an error: %v", i.(error))
			return
		}
	}
	val = results[0]
	return
}

var typeOfError = reflect.TypeOf((*error)(nil)).Elem()
var typeOfAny = reflect.TypeOf((*any)(nil)).Elem()

func checkFunc(fn reflect.Value, typ reflect.Type) error {
	if typ.Kind() != reflect.Func {
		return fmt.Errorf("value is a %s, not a func: %v", fn.Kind().String(), fn)
	}
	if typ.NumOut() < 1 {
		return fmt.Errorf("cannot inject function with no return values: %+v", fn)
	}
	if typ.NumOut() > 2 {
		return fmt.Errorf("cannot inject function with more than 2 return values: %+v", fn)
	}
	if typ.NumOut() == 2 && typ.Out(1) != typeOfError {
		return fmt.Errorf("cannot inject function with a non-error second return value: %s. %+v", typ.Out(1).String(), fn)
	}
	if typ.Out(0) == typeOfError {
		return fmt.Errorf("cannot inject function with only an error return value: %s. %+v", typ.Out(0).String(), fn)
	}
	return nil
}

func As[T any](v T) Injectable {
	t := reflect.TypeOf((*T)(nil)).Elem()
	return asValueInjectable{t, func(i Injector) (reflect.Value, error) {
		val, err := valueOf(reflect.ValueOf(v))
		if err != nil {
			return reflect.Value{}, err
		}
		return val.Inject(i)
	}}
}

type asValueInjectable struct {
	t reflect.Type
	f func(Injector) (reflect.Value, error)
}

func (v asValueInjectable) Type() reflect.Type                       { return v.t }
func (v asValueInjectable) Inject(i Injector) (reflect.Value, error) { return v.f(i) }

type valueInjectable reflect.Value

func (v valueInjectable) Type() reflect.Type                     { return reflect.Value(v).Type() }
func (v valueInjectable) Inject(Injector) (reflect.Value, error) { return reflect.Value(v), nil }

type funcInjectable reflect.Value

func (f funcInjectable) Type() reflect.Type                       { return reflect.Value(f).Type().Out(0) }
func (f funcInjectable) Inject(j Injector) (reflect.Value, error) { return inject(j, reflect.Value(f)) }

type sliceInjectable struct {
	typ     reflect.Type
	members []Injectable
}

func (s *sliceInjectable) Type() reflect.Type { return s.typ }
func (s *sliceInjectable) Inject(j Injector) (reflect.Value, error) {
	vs := reflect.MakeSlice(s.typ, 0, len(s.members))
	for _, m := range s.members {
		v, err := m.Inject(j)
		if err != nil {
			return reflect.Value{}, err
		}
		vs = reflect.Append(vs, v)
	}
	return vs, nil
}

type Scope struct {
	parent *Scope
	values map[reflect.Type]Injectable
}

func (ph *Scope) Scope(args ...any) (*Scope, error) {
	h := pool.Get().(*Scope)
	h.parent = ph
	h.add(As[Injector](h))
	err := h.add(args...)
	if err != nil {
		return nil, err
	}
	return h, nil
}

func (h *Scope) add(args ...any) error {
	for _, v := range args {
		if inj, ok := v.(Injectable); ok {
			if old, ok := h.values[inj.Type()]; ok {
				return fmt.Errorf("cannot add value to scope with the same type already present. old: %+v, new: %+v", old, v)
			} else {
				h.values[inj.Type()] = inj
				return nil
			}
		}
		val := reflect.ValueOf(v)
		switch val.Kind() {
		case reflect.Slice:
			// Only process unnamed slice types
			if val.Type().Name() == "" {
				if err := h.addSlice(val); err != nil {
					return err
				}
				continue
			}
			fallthrough
		default:
			if inj, err := valueOf(val); err != nil {
				return err
			} else if old, ok := h.values[inj.Type()]; ok {
				return fmt.Errorf("cannot add value to scope with the same type already present. old: %+v, new: %+v", old, val.Interface())
			} else {
				h.values[inj.Type()] = inj
			}
		}
	}
	return nil
}

func valueOf(v reflect.Value) (Injectable, error) {
	if v.Type() == typeOfAny {
		v = v.Elem()
	}
	typ := v.Type()
	switch v.Kind() {
	case reflect.Slice:
		return nil, fmt.Errorf("cannot convert value of type %v to injectable", typ.String())
	case reflect.Func:
		if err := checkFunc(v, typ); err != nil {
			return nil, err
		}
		return funcInjectable(v), nil
	default:
		return valueInjectable(v), nil
	}
}

func (h *Scope) addSlice(v reflect.Value) error {
	t := v.Type()
	te := t.Elem()
	if te == typeOfAny {
		if v.Len() == 0 {
			return nil // fmt.Errorf("cannot inject empty slice of any")
		}
		i, err := valueOf(v.Index(0))
		if err != nil {
			return err
		}
		te = i.Type()
		t = reflect.SliceOf(te)
	}
	var si *sliceInjectable
	if old, ok := h.values[t]; ok {
		// Case 1: It exists in this scope: use it
		if si, ok = old.(*sliceInjectable); !ok {
			panic(fmt.Errorf("type should be a sliceInjectable, got %+v (%T)", old, old))
		}
		si.members = slices.Grow(si.members, v.Len())
	} else if old, _ := h.parent.getInjectable(t); old != nil {
		// Case 2: It exists in ancestor scope: clone it
		if si, ok = old.(*sliceInjectable); !ok {
			panic(fmt.Errorf("type should be a sliceInjectable, got %+v (%T)", old, old))
		}
		// explicitly copy to avoid reallocation when appending
		mem := make([]Injectable, len(si.members), len(si.members)+v.Len())
		copy(mem, si.members)
		si = &sliceInjectable{typ: t, members: mem}
	} else {
		// Case 3: Doesn't exist yet: make a new one
		si = &sliceInjectable{typ: t, members: make([]Injectable, 0, v.Len())}
	}
	// Append to the selected object
	for i := 0; i < v.Len(); i++ {
		iv, err := valueOf(v.Index(i))
		if err != nil {
			return err
		}
		if iv.Type() != te {
			return fmt.Errorf("element not of same type as slice: %s != %s", iv, te)
		}
		si.members = append(si.members, iv)
	}
	h.values[t] = si
	return nil
}

func (h *Scope) Get(t reflect.Type) (reflect.Value, error) {
	i, _ := h.getInjectable(t)
	if i == nil {
		return reflect.Value{}, fmt.Errorf("failed to find type: %s", t.String())
	}
	return i.Inject(h)
}

func (h *Scope) getInjectable(t reflect.Type) (Injectable, int) {
	if h == nil {
		return nil, 0
	}
	if i, ok := h.values[t]; ok {
		return i, 1
	} else {
		i, n := h.parent.getInjectable(t)
		return i, n + 1
	}
}

var pool sync.Pool = sync.Pool{New: func() any {
	return &Scope{
		parent: nil,
		values: make(map[reflect.Type]Injectable),
	}
}}

func (h *Scope) PutPool() {
	h.parent = nil
	clear(h.values)
	pool.Put(h)
}