From 70cb79b23390abdc2340668fc823ab14f98984bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Patrick=20L=C3=BChne?= <patrick-github@luehne.de>
Date: Mon, 6 Mar 2017 15:40:23 +0100
Subject: [PATCH] Implemented choice rules.

---
 include/anthem/Context.h          |  4 ++++
 include/anthem/Head.h             | 39 +++++++++++++++++++++++++++----
 include/anthem/StatementVisitor.h | 19 ++++++++++++++-
 3 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/include/anthem/Context.h b/include/anthem/Context.h
index a488b5b..0149a20 100644
--- a/include/anthem/Context.h
+++ b/include/anthem/Context.h
@@ -19,12 +19,16 @@ struct Context
 	void reset()
 	{
 		headTerms.clear();
+		isChoiceRule = false;
+		numberOfHeadLiterals = 0;
 		auxiliaryBodyLiteralID = 1;
 	}
 
 	output::Logger logger;
 
 	std::vector<const Clingo::AST::Term *> headTerms;
+	bool isChoiceRule = false;
+	size_t numberOfHeadLiterals = 0;
 	size_t auxiliaryBodyLiteralID = 1;
 };
 
diff --git a/include/anthem/Head.h b/include/anthem/Head.h
index 3211a9d..b695657 100644
--- a/include/anthem/Head.h
+++ b/include/anthem/Head.h
@@ -86,15 +86,20 @@ struct LiteralCollectFunctionTermsVisitor
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
+// TODO: rename, because not only terms are collected anymore
 struct HeadLiteralCollectFunctionTermsVisitor
 {
 	void visit(const Clingo::AST::Literal &literal, const Clingo::AST::HeadLiteral &, Context &context)
 	{
+		context.numberOfHeadLiterals = 1;
+
 		literal.data.accept(LiteralCollectFunctionTermsVisitor(), literal, context);
 	}
 
 	void visit(const Clingo::AST::Disjunction &disjunction, const Clingo::AST::HeadLiteral &headLiteral, Context &context)
 	{
+		context.numberOfHeadLiterals = disjunction.elements.size();
+
 		for (const auto &conditionalLiteral : disjunction.elements)
 		{
 			if (!conditionalLiteral.condition.empty())
@@ -104,9 +109,21 @@ struct HeadLiteralCollectFunctionTermsVisitor
 		}
 	}
 
-	void visit(const Clingo::AST::Aggregate &, const Clingo::AST::HeadLiteral &headLiteral, Context &context)
+	void visit(const Clingo::AST::Aggregate &aggregate, const Clingo::AST::HeadLiteral &headLiteral, Context &context)
 	{
-		throwErrorAtLocation(headLiteral.location, "“aggregate” head literals currently unsupported", context);
+		context.isChoiceRule = true;
+		context.numberOfHeadLiterals = aggregate.elements.size();
+
+		if (aggregate.left_guard || aggregate.right_guard)
+			throwErrorAtLocation(headLiteral.location, "aggregates with left or right guards not allowed", context);
+
+		for (const auto &conditionalLiteral : aggregate.elements)
+		{
+			if (!conditionalLiteral.condition.empty())
+				throwErrorAtLocation(headLiteral.location, "conditional literals in aggregates currently unsupported", context);
+
+			conditionalLiteral.literal.data.accept(LiteralCollectFunctionTermsVisitor(), conditionalLiteral.literal, context);
+		}
 	}
 
 	void visit(const Clingo::AST::HeadAggregate &, const Clingo::AST::HeadLiteral &headLiteral, Context &context)
@@ -244,9 +261,23 @@ struct HeadLiteralPrintSubstitutedVisitor
 		}
 	}
 
-	void visit(const Clingo::AST::Aggregate &, const Clingo::AST::HeadLiteral &headLiteral, Context &context)
+	void visit(const Clingo::AST::Aggregate &aggregate, const Clingo::AST::HeadLiteral &headLiteral, Context &context)
 	{
-		throwErrorAtLocation(headLiteral.location, "“aggregate” head literals currently unsupported", context);
+		if (aggregate.left_guard || aggregate.right_guard)
+			throwErrorAtLocation(headLiteral.location, "aggregates with left or right guards not allowed", context);
+
+		for (auto i = aggregate.elements.cbegin(); i != aggregate.elements.cend(); i++)
+		{
+			const auto &conditionalLiteral = *i;
+
+			if (!conditionalLiteral.condition.empty())
+				throwErrorAtLocation(headLiteral.location, "conditional head literals currently unsupported", context);
+
+			if (i != aggregate.elements.cbegin())
+				context.logger.outputStream() << " " << Clingo::AST::BinaryOperator::Or << " ";
+
+			visit(conditionalLiteral.literal, headLiteral, context);
+		}
 	}
 
 	void visit(const Clingo::AST::HeadAggregate &, const Clingo::AST::HeadLiteral &headLiteral, Context &context)
diff --git a/include/anthem/StatementVisitor.h b/include/anthem/StatementVisitor.h
index cb5bdab..ba5abd0 100644
--- a/include/anthem/StatementVisitor.h
+++ b/include/anthem/StatementVisitor.h
@@ -53,7 +53,7 @@ struct StatementVisitor
 			}
 		}
 
-		if (rule.body.empty() && context.headTerms.empty())
+		if (rule.body.empty() && context.headTerms.empty() && !context.isChoiceRule)
 			outputStream << Clingo::AST::Boolean({true});
 		else
 		{
@@ -72,6 +72,23 @@ struct StatementVisitor
 			}
 		}
 
+		// Handle choice rules
+		if (context.isChoiceRule)
+		{
+			const bool isFirstOutput = rule.body.empty() && context.headTerms.empty();
+
+			if (!isFirstOutput)
+				outputStream << " " << Clingo::AST::BinaryOperator::And << " ";
+
+			if (context.numberOfHeadLiterals > 1 && !isFirstOutput)
+				outputStream << "(";
+
+			rule.head.data.accept(HeadLiteralPrintSubstitutedVisitor(), rule.head, context);
+
+			if (context.numberOfHeadLiterals > 1 && !isFirstOutput)
+				outputStream << ")";
+		}
+
 		outputStream << " " << output::Operator("->") << " ";
 
 		// Print consequent of the implication