1 files changed, 23 insertions, 9 deletions

diff --git a/src/jcgp/backend/function/FunctionSet.java b/src/jcgp/backend/function/FunctionSet.java
index 4470ac8..78801fc 100644
--- a/src/jcgp/backend/function/FunctionSet.java
+++ b/src/jcgp/backend/function/FunctionSet.java
@@ -1,6 +1,7 @@
 package jcgp.backend.function;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.Iterator;
 
 
@@ -11,7 +12,7 @@ import java.util.Iterator;
  */
 public abstract class FunctionSet {
 		protected Function[] functionList;
-		protected ArrayList<Function> allowedFunctions;
+		protected ArrayList<Integer> allowedFunctions;
 		protected String name;
 		
 		public int getAllowedFunctionCount() {
@@ -23,7 +24,7 @@ public abstract class FunctionSet {
 		}
 		
 		public Function getAllowedFunction(int index) {
-			return allowedFunctions.get(index);
+			return functionList[allowedFunctions.get(index)];
 		}
 		
 		public Function getFunction(int index) {
@@ -32,7 +33,7 @@ public abstract class FunctionSet {
 		
 		public int getMaxArity(){
 			int arity = 0;
-			for (Function function : allowedFunctions) {
+			for (Function function : functionList) {
 				if (function.getArity() > arity) {
 					arity = function.getArity();
 				}
@@ -46,9 +47,9 @@ public abstract class FunctionSet {
 		
 		public void disableFunction(int index) {
 			if (index < functionList.length) {
-				for (Iterator<Function> iterator = allowedFunctions.iterator(); iterator.hasNext();) {
-					Function function = (Function) iterator.next();
-					if (function == functionList[index]) {
+				for (Iterator<Integer> iterator = allowedFunctions.iterator(); iterator.hasNext();) {
+					int function = iterator.next();
+					if (function == index) {
 						iterator.remove();
 					}
 				}
@@ -58,8 +59,9 @@ public abstract class FunctionSet {
 		}
 		
 		public void enableFunction(int index) {
-			if (!allowedFunctions.contains(functionList[index])) {
-				allowedFunctions.add(functionList[index]);
+			if (!allowedFunctions.contains(index)) {
+				allowedFunctions.add(index);
+				Collections.sort(allowedFunctions);
 			}
 		}
 		
@@ -69,6 +71,18 @@ public abstract class FunctionSet {
 		}
 		
 		public boolean isEnabled(Function f) {
-			return allowedFunctions.contains(f);
+			for (int i = 0; i < allowedFunctions.size(); i++) {
+				if (functionList[allowedFunctions.get(i)] == f) {
+					return true;
+				}
+			}
+			return false;
+		}
+		
+		protected void enableAll() {
+			allowedFunctions = new ArrayList<Integer>();
+			for (int i = 0; i < functionList.length; i++) {
+				allowedFunctions.add(i);
+			}
 		}
 	}
\ No newline at end of file