Example of a trealla-prolog/go native predicate returning multiple choice points

channelpredicate.go

package trealla_test

import (
	"context"
	"sync/atomic"
	"testing"

	"sync"

	"github.com/trealla-prolog/go/trealla"
	"github.com/trealla-prolog/go/trealla/terms"
)

type ChannelReceiver struct {
	m  *sync.Map
	id *atomic.Int64
	fn func(goal trealla.Term) chan trealla.Term

	// functors
	start trealla.Atom
	recv  trealla.Atom
}

func NewChannelReceiver(functor trealla.Atom, fn func(goal trealla.Term) chan trealla.Term) *ChannelReceiver {
	return &ChannelReceiver{
		m:     new(sync.Map),
		id:    new(atomic.Int64),
		fn:    fn,
		start: functor,
		recv:  "$" + functor + "_recv",
	}
}

func (cr *ChannelReceiver) channel_start_2(_ trealla.Prolog, _ trealla.Subquery, goal trealla.Term) trealla.Term {
	goalcmp, ok := goal.(trealla.Compound)
	if !ok {
		return terms.Throw(terms.TypeError("compound", goalcmp, terms.PI(goal)))
	}
	input := goalcmp.Args[0]
	result := goalcmp.Args[1]

	id := cr.id.Add(1)
	ch := cr.fn(input)
	cr.m.Store(id, ch)
	return trealla.Atom("call").Of(cr.recv.Of(id, result))
}

func (cr *ChannelReceiver) channel_next_2(_ trealla.Prolog, _ trealla.Subquery, goal trealla.Term) trealla.Term {
	goalcmp, ok := goal.(trealla.Compound)
	if !ok {
		return terms.Throw(terms.TypeError("compound", goalcmp, terms.PI(goal)))
	}
	idTerm := goalcmp.Args[0]
	id, ok := idTerm.(int64)
	if !ok {
		return terms.Throw(terms.TypeError("integer", idTerm, terms.PI(goal)))
	}
	result := goalcmp.Args[1]

	ch, ok := cr.m.Load(id)
	if !ok {
		return trealla.Atom("fail")
	}
	got, ok := <-ch.(chan trealla.Term)
	if !ok {
		cr.m.Delete(id)
		return trealla.Atom("fail")
	}
	if cmp, ok := got.(trealla.Compound); ok && cmp.Functor == "throw" {
		cr.m.Delete(id)
		return cmp
	}
	return trealla.Atom("call").Of(trealla.Atom(";").Of(trealla.Atom("=").Of(result, got), cr.recv.Of(id, result)))
}

func (cr *ChannelReceiver) Register(ctx context.Context, pl trealla.Prolog) error {
	if err := pl.Register(ctx, string(cr.start), 2, cr.channel_start_2); err != nil {
		return err
	}
	if err := pl.Register(ctx, string(cr.recv), 2, cr.channel_next_2); err != nil {
		return err
	}
	return nil
}

func TestInteropChoice(t *testing.T) {
	fn := func(input trealla.Term) chan trealla.Term {
		ch := make(chan trealla.Term)
		count, ok := input.(int64)
		if !ok {
			go func() {
				ch <- terms.Throw(terms.TypeError("integer", input, trealla.Atom("/").Of("countdown", 2)))
				close(ch)
			}()
			return ch
		}
		go func() {
			for i := int64(0); i < count; i++ {
				ch <- trealla.Term(i)
			}
			close(ch)
		}()
		return ch
	}
	cr := NewChannelReceiver("countdown", fn)
	pl, err := trealla.New()
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()
	cr.Register(ctx, pl)

	t.Run("success", func(t *testing.T) {
		q := pl.Query(ctx, "countdown(10, X)")
		var n int64
		for q.Next(ctx) {
			t.Logf("got: %v", q.Current())
			if q.Current().Solution["X"] != n {
				t.Error("unexpected solution:", q.Current())
			}
			n++
		}
		if q.Err() != nil {
			t.Error(q.Err())
		}
	})

	t.Run("bad arg", func(t *testing.T) {
		q := pl.Query(ctx, "countdown(foobar, X)")
		for q.Next(ctx) {
		}
		if q.Err() == nil {
			t.Error("expected error")
		}
	})
}