go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/projects/nodes/pkg/incrutil/map2.go

/*

Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.

*/

package incrutil

import (
	"context"
	"fmt"

	"github.com/wcharczuk/go-incr"
	"go.charczuk.com/sdk/errutil"
)

// MapN applies a function to given list of input incrementals and returns
// a new incremental of the output type of that function.
func Map2[A, B, C any](scope incr.Scope, fn Map2Func[A, B, C]) Map2Incr[A, B, C] {
	return incr.WithinScope(scope, &map2Incr[A, B, C]{
		n:  incr.NewNode("map2"),
		fn: fn,
	})
}

// Map2Func is the function that the ApplyN incremental applies.
type Map2Func[A, B, C any] func(context.Context, A, B) (C, error)

// MapNIncr is a type of incremental that can add inputs over time.
type Map2Incr[A, B, C any] interface {
	incr.Incr[C]
	ISetInput
	IRemoveInput
	IInputs
	IRawFunction
	ISetRawFunction
	IRawValue
	ISetRawValue
}

var (
	_ incr.Incr[bool]             = (*map2Incr[int, string, bool])(nil)
	_ Map2Incr[int, string, bool] = (*map2Incr[int, string, bool])(nil)
	_ incr.INode                  = (*map2Incr[int, string, bool])(nil)
	_ incr.IStabilize             = (*map2Incr[int, string, bool])(nil)
	_ fmt.Stringer                = (*map2Incr[int, string, bool])(nil)
)

type map2Incr[A, B, C any] struct {
	n   *incr.Node
	a   incr.Incr[A]
	b   incr.Incr[B]
	fn  Map2Func[A, B, C]
	val C
}

func (m *map2Incr[A, B, C]) Parents() (out []incr.INode) {
	if m.a != nil {
		out = append(out, m.a)
	}
	if m.b != nil {
		out = append(out, m.b)
	}
	return
}

func (m *map2Incr[A, B, C]) RawValue() any {
	return m.val
}

func (m *map2Incr[A, B, C]) SetRawValue(value any) error {
	var typed C
	if err := CastAny(value, &typed); err != nil {
		return err
	}
	m.val = typed
	return nil
}

func (mn *map2Incr[A, B, C]) RawFunction() any {
	return mn.fn
}

func (mn *map2Incr[A, B, C]) Inputs() map[string]incr.INode {
	output := make(map[string]incr.INode)
	if mn.a != nil {
		output["a"] = mn.a
	}
	if mn.b != nil {
		output["b"] = mn.b
	}
	return output
}

func (mn *map2Incr[A, B, C]) SetInput(name string, input incr.INode, _ bool) error {
	if name == "0" || name == "a" {
		typed, ok := input.(incr.Incr[A])
		if !ok {
			return fmt.Errorf("set input; invalid input type for input %v: %T", name, input)
		}
		mn.a = typed
		LinkNodes(mn, mn.a)
		return nil
	} else if name == "1" || name == "b" {
		typed, ok := input.(incr.Incr[B])
		if !ok {
			return fmt.Errorf("set input; invalid input type for input %v: %T", name, input)
		}
		mn.b = typed
		LinkNodes(mn, mn.b)
		return nil
	}
	return fmt.Errorf("set input; invalid input name: %v", name)
}

func (mn *map2Incr[A, B, C]) RemoveInput(name string) error {
	if name == "0" || name == "a" {
		if mn.a != nil {
			UnlinkNodes(mn, mn.a)
			mn.a = nil
			return nil
		}
		return fmt.Errorf("remove input; input unset: %v", name)
	} else if name == "1" || name == "b" {
		if mn.b != nil {
			UnlinkNodes(mn, mn.b)
			mn.b = nil
			return nil
		}
		return fmt.Errorf("remove input; input unset: %v", name)
	}
	return fmt.Errorf("remove input; invalid input name: %v", name)
}

func (mn *map2Incr[A, B, C]) SetRawFunction(fn any) error {
	typed, ok := fn.(Map2Func[A, B, C])
	if !ok {
		return fmt.Errorf("invalid function type for node: %T", fn)
	}
	mn.fn = typed
	return nil
}

func (m *map2Incr[A, B, C]) Node() *incr.Node { return m.n }

func (m *map2Incr[A, B, C]) Value() C { return m.val }

func (m *map2Incr[A, B, C]) Stabilize(ctx context.Context) (err error) {
	var val C
	var mva A
	var mvb B
	if m.a != nil {
		mva = m.a.Value()
	}
	if m.b != nil {
		mvb = m.b.Value()
	}
	val, err = m.fn(ctx, mva, mvb)
	if err != nil {
		err = errutil.New(fmt.Errorf("%v; %w", m, err))
		return
	}
	m.val = val
	return nil
}

func (m *map2Incr[A, B, C]) String() string {
	return m.n.String()
}