From 85a44f15057dc53e85e3287ceea362322b77fa82 Mon Sep 17 00:00:00 2001
From: Renaud Hartert <91352809+rhartert@users.noreply.github.com>
Date: Mon, 27 May 2024 19:16:06 +0900
Subject: [PATCH] Add predicate parsing

---
 fzn/fzn_test.go   | 84 +++++++++++++++++++++++++++++++++++++++++++++--
 fzn/predicates.go | 75 +++++++++++++++++++++++++++++++++++++-----
 fzn/types.go      | 10 +++++-
 3 files changed, 158 insertions(+), 11 deletions(-)

diff --git a/fzn/fzn_test.go b/fzn/fzn_test.go
index ffd83ed..960e596 100644
--- a/fzn/fzn_test.go
+++ b/fzn/fzn_test.go
@@ -93,9 +93,89 @@ func TestParse_comment(t *testing.T) {
 func TestParse_predicate(t *testing.T) {
 	testParse(t, []testCase{
 		{
-			input: "predicate foo(int)",
+			input:   "predicate;",
+			wantErr: true,
+		},
+		{
+			input:   "predicate foo;",
+			wantErr: true,
+		},
+		{
+			input:   "predicate foo(int)",
+			wantErr: true,
+		},
+		{
+			input:   "predicate (int);",
+			wantErr: true,
+		},
+		{
+			input:   "predicate foo int : X;",
+			wantErr: true,
+		},
+		{
+			input:   "predicate (int : X);",
+			wantErr: true,
+		},
+		{
+			input:   "predicate foo(set : X);",
+			wantErr: true,
+		},
+		{
+			input:   "predicate foo(var var : X);",
+			wantErr: true,
+		},
+		{
+			input:   "predicate foo(array of int : X);",
+			wantErr: true,
+		},
+		{
+			input:   "predicate foo(int: A)",
+			wantErr: true,
+		},
+		{
+			input:   "predicate foo(int: int)",
+			wantErr: true,
+		},
+		{
+			input:   "predicate foo(bool: A int: B float: C);",
+			wantErr: true,
+		},
+		{
+			input: "predicate foo(int: A);",
+			want: instruction{
+				Predicate: &Predicate{
+					Identifier: "foo",
+					Parameters: []PredParam{
+						{Identifier: "A", ParType: ParTypeInt},
+					},
+				},
+			},
+		},
+		{
+			input: "predicate foo(bool: A, int: B, float: C, var int: X);",
 			want: instruction{
-				Predicate: &Predicate{Value: "predicate foo ( int )"},
+				Predicate: &Predicate{
+					Identifier: "foo",
+					Parameters: []PredParam{
+						{Identifier: "A", ParType: ParTypeBool},
+						{Identifier: "B", ParType: ParTypeInt},
+						{Identifier: "C", ParType: ParTypeFloat},
+						{Identifier: "X", VarType: VarTypeIntRange},
+					},
+				},
+			},
+		},
+		{
+			input: "predicate foo(array [1..10] of var float: X);",
+			want: instruction{
+				Predicate: &Predicate{
+					Identifier: "foo",
+					Parameters: []PredParam{{
+						Identifier: "X",
+						Array:      &Array{1, 10},
+						VarType:    VarTypeFloatRange,
+					}},
+				},
 			},
 		},
 	})
diff --git a/fzn/predicates.go b/fzn/predicates.go
index ee79de0..fe3ee54 100644
--- a/fzn/predicates.go
+++ b/fzn/predicates.go
@@ -2,7 +2,6 @@ package fzn
 
 import (
 	"fmt"
-	"strings"
 
 	"github.com/rhartert/gofzn/fzn/tok"
 )
@@ -14,15 +13,75 @@ func isPredicate(p *parser) bool {
 // parsePredicate returns a string of all the token value contained in the
 // predicate. Said otherwise, the parser effectively treats predicates as
 // comments
-func parsePredicate(p *parser) (*Predicate, error) {
+func parsePredicate(p *parser) (pred *Predicate, err error) {
 	if p.next().Type != tok.Predicate {
 		return nil, fmt.Errorf("not a predicate")
 	}
-	sb := strings.Builder{}
-	sb.WriteString("predicate")
-	for t := p.next(); t.Type != tok.EOF; t = p.next() {
-		sb.WriteByte(' ')
-		sb.WriteString(t.Value)
+
+	pred = &Predicate{}
+	pred.Identifier, err = parseIdentifier(p)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing predicate identifier: %w", err)
+	}
+
+	if !p.nextIf(tok.TupleStart) {
+		return nil, fmt.Errorf("missing ( after predicate identifier")
+	}
+
+	for !p.nextIf(tok.TupleEnd) {
+		pp, err := parsePredicateParam(p)
+		if err != nil {
+			return nil, fmt.Errorf("error parsing predicate parameter: %w", err)
+		}
+		pred.Parameters = append(pred.Parameters, pp)
+
+		if !p.nextIf(tok.Comma) && p.lookAhead(0).Type != tok.TupleEnd {
+			return nil, fmt.Errorf("missing comma")
+		}
+	}
+
+	if !p.nextIf(tok.EOI) {
+		return nil, fmt.Errorf("missing end of predicate instruction")
+	}
+
+	return pred, nil
+}
+
+func parsePredicateParam(p *parser) (PredParam, error) {
+	pp := PredParam{}
+
+	if p.lookAhead(0).Type == tok.Array {
+		a, err := parseArrayOf(p)
+		if err != nil {
+			return PredParam{}, err
+		}
+		pp.Array = a
+	}
+
+	switch p.lookAhead(0).Type {
+	case tok.Var:
+		v, err := parseVariable(p)
+		if err != nil {
+			return PredParam{}, err
+		}
+		pp.VarType = v.Type
+	default: // parameter
+		pt, err := parseParType(p)
+		if err != nil {
+			return PredParam{}, err
+		}
+		pp.ParType = pt
 	}
-	return &Predicate{Value: sb.String()}, nil
+
+	if !p.nextIf(tok.Colon) {
+		return PredParam{}, fmt.Errorf("missing ':'")
+	}
+
+	id, err := parseIdentifier(p)
+	if err != nil {
+		return PredParam{}, err
+	}
+	pp.Identifier = id
+
+	return pp, nil
 }
diff --git a/fzn/types.go b/fzn/types.go
index eb08a00..bb99d0c 100644
--- a/fzn/types.go
+++ b/fzn/types.go
@@ -1,7 +1,15 @@
 package fzn
 
 type Predicate struct {
-	Value string
+	Identifier string
+	Parameters []PredParam
+}
+
+type PredParam struct {
+	Identifier string
+	Array      *Array
+	VarType    VarType
+	ParType    ParType
 }
 
 type ParamDeclaration struct {