Additional cleanup and refactoring in ContextFunctionCatalog

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/config/ContextFunctionCatalogAutoConfiguration.java

@@ -115,7 +115,7 @@ public class ContextFunctionCatalogAutoConfiguration {
 
 		@Override
 		public FunctionRegistration<?> getRegistration(Object function) {
-			return this.processor.getRegistration(function);
+			return function == null ? null : this.processor.getRegistration(function);
 		}
 
 		@Override
@@ -125,6 +125,27 @@ public class ContextFunctionCatalogAutoConfiguration {
 			this.processor.register(registration);
 		}
 
+		@Override
+		public Set<String> getNames(Class<?> type) {
+			if (Supplier.class.isAssignableFrom(type)) {
+				return this.processor.suppliers.keySet();
+			}
+			if (Consumer.class.isAssignableFrom(type)) {
+				return this.processor.consumers.keySet();
+			}
+			if (Function.class.isAssignableFrom(type)) {
+				return this.processor.functions.keySet();
+			}
+			return Collections.emptySet();
+		}
+
+		@Override
+		public int size() {
+			return this.processor.suppliers.size()
+					+ this.processor.functions.size()
+					+ this.processor.consumers.size();
+		}
+
 		@Override
 		@SuppressWarnings("unchecked")
 		protected <T> T doLookup(Class<?> type, String name) {
@@ -150,55 +171,6 @@ public class ContextFunctionCatalogAutoConfiguration {
 			return function;
 		}
 
-		@Override
-		public Set<String> getNames(Class<?> type) {
-			if (Supplier.class.isAssignableFrom(type)) {
-				return this.processor.getSuppliers();
-			}
-			if (Consumer.class.isAssignableFrom(type)) {
-				return this.processor.getConsumers();
-			}
-			if (Function.class.isAssignableFrom(type)) {
-				return this.processor.getFunctions();
-			}
-			return Collections.emptySet();
-		}
-
-		@Override
-		public int size() {
-			return this.processor.getSuppliers().size()
-					+ this.processor.getFunctions().size()
-					+ this.processor.getConsumers().size();
-		}
-
-	}
-
-	@Configuration
-	@ConditionalOnClass(Gson.class)
-	@ConditionalOnBean(Gson.class)
-	@Conditional(PreferGsonOrMissingJacksonCondition.class)
-	protected static class GsonConfiguration {
-
-		@Bean
-		public GsonMapper jsonMapper(Gson gson) {
-			return new GsonMapper(gson);
-		}
-
-	}
-
-	@Configuration
-	@ConditionalOnClass(ObjectMapper.class)
-	@ConditionalOnBean(ObjectMapper.class)
-	// @checkstyle:off
-	@ConditionalOnProperty(name = ContextFunctionCatalogAutoConfiguration.PREFERRED_MAPPER_PROPERTY, havingValue = "jackson", matchIfMissing = true)
-	// @checkstyle:on
-	protected static class JacksonConfiguration {
-
-		@Bean
-		public JacksonMapper jsonMapper(ObjectMapper mapper) {
-			return new JacksonMapper(mapper);
-		}
-
 	}
 
 	@Component
@@ -207,22 +179,20 @@ public class ContextFunctionCatalogAutoConfiguration {
 
 		private Log logger = LogFactory.getLog(ContextFunctionRegistry.class);
 
+		private ApplicationEventPublisher applicationEventPublisher;
+
+		private ConfigurableListableBeanFactory beanFactory;
+
 		private Map<String, Object> suppliers = new ConcurrentHashMap<>();
 
 		private Map<String, Object> functions = new ConcurrentHashMap<>();
 
 		private Map<String, Object> consumers = new ConcurrentHashMap<>();
 
-		private ApplicationEventPublisher applicationEventPublisher;
-
-		private ConfigurableListableBeanFactory beanFactory;
-
 		private Map<Object, String> names = new ConcurrentHashMap<>();
 
 		private Map<String, FunctionType> types = new ConcurrentHashMap<>();
 
-		private Map<String, Object> allFunctions = new ConcurrentHashMap<>();
-
 		/**
 		 * Will collect all suppliers, functions, consumers and function registration as
 		 * late as possible in the lifecycle.
@@ -273,48 +243,34 @@ public class ContextFunctionCatalogAutoConfiguration {
 		}
 
 		FunctionRegistration<?> getRegistration(Object function) {
-			if (function == null || !this.names.containsKey(function)) {
+			if (names.containsKey(function)) {
-				return null;
+				return new FunctionRegistration<>(function, this.names.get(function))
+						.type(findType(function).getType());
 			}
-			return new FunctionRegistration<>(function, this.names.get(function))
+			return null;
-					.type(findType(function).getType());
 		}
 
 		Supplier<?> lookupSupplier(String name) {
-			Object function = compose(name);
+			return (Supplier<?>) lookup(name, this.suppliers, Supplier.class);
-			if (function instanceof Supplier) {
-				return (Supplier<?>) function;
-			}
-			else {
-				logger.warn("The resulting composition is is of type "
-						+ types.get(normalizeName(name))
-						+ " and can not be cast to Supplier");
-			}
-			return null;
 		}
 
 		Function<?, ?> lookupFunction(String name) {
-			Object function = compose(name);
+			return (Function<?, ?>) lookup(name, this.functions, Function.class);
-			if (function instanceof Function) {
-				return (Function<?, ?>) function;
-			}
-			else {
-				logger.warn("The resulting composition is is of type "
-						+ types.get(normalizeName(name))
-						+ " and can not be cast to Function");
-			}
-			return null;
 		}
 
 		Consumer<?> lookupConsumer(String name) {
-			Object function = compose(name);
+			return (Consumer<?>) lookup(name, this.consumers, Consumer.class);
-			if (function instanceof Consumer) {
+		}
-				return (Consumer<?>) function;
+
+		@SuppressWarnings("unchecked")
+		private Object lookup(String name, @SuppressWarnings("rawtypes") Map lookup, Class<?> typeOfFunction) {
+			Object function = compose(name, lookup);
+			if (function  != null && typeOfFunction.isAssignableFrom(function.getClass())) {
+				return function;
 			}
 			else {
 				logger.warn("The resulting composition is is of type "
-						+ types.get(normalizeName(name))
+						+ types.get(normalizeName(name)));
-						+ " and can not be cast to Consumer");
 			}
 			return null;
 		}
@@ -337,39 +293,22 @@ public class ContextFunctionCatalogAutoConfiguration {
 			return name.replaceAll(",", "|").trim();
 		}
 
-		private Set<String> getSuppliers() {
+		private Object compose(String name, Map<String, Object> lookup) {
-			return this.suppliers.keySet();
-		}
-
-		private Set<String> getFunctions() {
-			return this.functions.keySet();
-		}
-
-		private Set<String> getConsumers() {
-			return this.consumers.keySet();
-		}
-
-		private Object compose(String name) {
 			name = normalizeName(name);
 			Object composedFunction = null;
-			if (allFunctions.containsKey(name)) {
+			if (lookup.containsKey(name)) {
-				composedFunction = allFunctions.get(name);
+				composedFunction = lookup.get(name);
 			}
 			else {
-				/*
+				if (name.equals("") && lookup.size() == 1) {
-				 * Need to revisit since "" implies a single function or supplier or... By
+					composedFunction = lookup.values().iterator().next();
-				 * using combined map we may have a single function but multiple consumers
-				 * scenario
-				 */
-				if (name.equals("") && allFunctions.size() == 1) {
-					composedFunction = allFunctions.values().iterator().next();
 				}
 				else {
 					String[] stages = StringUtils.delimitedListToStringArray(name, "|");
 					if (Stream.of(stages)
-							.allMatch(funcName -> allFunctions.containsKey(funcName))) {
+							.allMatch(funcName -> contains(funcName))) {
 						List<Object> composableFunctions = Stream.of(stages)
-								.map(funcName -> allFunctions.get(funcName))
+								.map(funcName -> find(funcName))
 								.collect(Collectors.toList());
 						composedFunction = composableFunctions.stream()
 								.reduce((a, z) -> composeFunctions(a, z))
@@ -382,7 +321,7 @@ public class ContextFunctionCatalogAutoConfiguration {
 									.get(stages[stages.length - 1]);
 							this.types.put(name, FunctionType.compose(input, output));
 							this.names.put(composedFunction, name);
-							allFunctions.put(name, composedFunction);
+							lookup.put(name, composedFunction);
 						}
 					}
 				}
@@ -390,6 +329,21 @@ public class ContextFunctionCatalogAutoConfiguration {
 			return composedFunction;
 		}
 
+		private boolean contains(String name) {
+			return suppliers.containsKey(name) || functions.containsKey(name) || consumers.containsKey(name);
+		}
+
+		private Object find(String name) {
+			Object result = suppliers.get(name);
+			if (result == null) {
+				result = functions.get(name);
+			}
+			if (result == null) {
+				result = consumers.get(name);
+			}
+			return result;
+		}
+
 		@SuppressWarnings("unchecked")
 		private Object composeFunctions(Object a, Object b) {
 			if (a instanceof Supplier && b instanceof Function) {
@@ -471,21 +425,18 @@ public class ContextFunctionCatalogAutoConfiguration {
 				type = Supplier.class;
 				for (String name : registration.getNames()) {
 					this.suppliers.put(name, registration.getTarget());
-					this.allFunctions.put(name, registration.getTarget());
 				}
 			}
 			else if (target instanceof Consumer) {
 				type = Consumer.class;
 				for (String name : registration.getNames()) {
 					this.consumers.put(name, registration.getTarget());
-					this.allFunctions.put(name, registration.getTarget());
 				}
 			}
 			else if (target instanceof Function) {
 				type = Function.class;
 				for (String name : registration.getNames()) {
 					this.functions.put(name, registration.getTarget());
-					this.allFunctions.put(name, registration.getTarget());
 				}
 			}
 			else {
@@ -539,21 +490,18 @@ public class ContextFunctionCatalogAutoConfiguration {
 			Object target = registration.getTarget();
 			boolean isolated = getClass().getClassLoader() != target.getClass()
 					.getClassLoader();
-			if (target instanceof Supplier<?>) {
+			if (isolated) {
-				if (isolated) {
+				if (target instanceof Supplier<?> && isolated) {
 					target = new IsolatedSupplier((Supplier<?>) target);
 				}
-			}
+				else if (target instanceof Function<?, ?>) {
-			else if (target instanceof Function<?, ?>) {
-				if (isolated) {
 					target = new IsolatedFunction((Function<?, ?>) target);
 				}
-			}
+				else if (target instanceof Consumer<?>) {
-			else if (target instanceof Consumer<?>) {
-				if (isolated) {
 					target = new IsolatedConsumer((Consumer<?>) target);
 				}
 			}
+
 			registration.target(target);
 			return registration;
 		}
@@ -668,4 +616,33 @@ public class ContextFunctionCatalogAutoConfiguration {
 
 	}
 
+	@Configuration
+	@ConditionalOnClass(Gson.class)
+	@ConditionalOnBean(Gson.class)
+	@Conditional(PreferGsonOrMissingJacksonCondition.class)
+	protected static class GsonConfiguration {
+
+		@Bean
+		public GsonMapper jsonMapper(Gson gson) {
+			return new GsonMapper(gson);
+		}
+
+	}
+
+	@Configuration
+	@ConditionalOnClass(ObjectMapper.class)
+	@ConditionalOnBean(ObjectMapper.class)
+	// @checkstyle:off
+	@ConditionalOnProperty(name = ContextFunctionCatalogAutoConfiguration.PREFERRED_MAPPER_PROPERTY, havingValue = "jackson", matchIfMissing = true)
+	// @checkstyle:on
+	protected static class JacksonConfiguration {
+
+		@Bean
+		public JacksonMapper jsonMapper(ObjectMapper mapper) {
+			return new JacksonMapper(mapper);
+		}
+
+	}
+
+
 }