Add support for composing function of Message with plain function

Fixes gh-267 at least for the most common cases.

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/AbstractComposableFunctionRegistry.java

@@ -362,6 +362,12 @@ public abstract class AbstractComposableFunctionRegistry implements FunctionRegi
 		FunctionType bType = bReg.getType();
 		Object a = aReg.getTarget();
 		Object b = bReg.getTarget();
+		if (aType != null && bType != null) {
+			if (aType.isMessage() && !bType.isMessage()) {
+				bType = bType.message();
+				b = message(b);
+			}
+		}
 		Object composedFunction = null;
 		if (a instanceof Supplier && b instanceof Function) {
 			Supplier<Flux<Object>> supplier = (Supplier<Flux<Object>>) a;
@@ -420,6 +426,19 @@ public abstract class AbstractComposableFunctionRegistry implements FunctionRegi
 				.type(FunctionType.compose(aType, bType));
 	}
 
+	private Object message(Object input) {
+		if (input instanceof Supplier) {
+			return new MessageSupplier((Supplier<?>) input);
+		}
+		if (input instanceof Consumer) {
+			return new MessageConsumer((Consumer<?>) input);
+		}
+		if (input instanceof Function) {
+			return new MessageFunction((Function<?, ?>) input);
+		}
+		return input;
+	}
+
 	@SuppressWarnings("unchecked")
 	private <T> T doLookup(Class<?> type, String name) {
 		T function = null;

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/MessageConsumer.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright 2019-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.function.context.catalog;
+
+import java.util.function.Consumer;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+
+import org.springframework.messaging.Message;
+
+/**
+ * @author Dave Syer
+ */
+public class MessageConsumer implements Consumer<Publisher<Message<?>>> {
+
+	private Consumer<Object> delegate;
+
+	@SuppressWarnings("unchecked")
+	public MessageConsumer(Consumer<?> input) {
+		this.delegate = (Consumer<Object>) input;
+	}
+
+	@Override
+	public void accept(Publisher<Message<?>> input) {
+		Flux.from(input).map(Message::getPayload).subscribe(this.delegate::accept);
+	}
+
+}

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/MessageFunction.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2019-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.function.context.catalog;
+
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import org.springframework.cloud.function.core.FluxConsumer;
+import org.springframework.cloud.function.core.FluxFunction;
+import org.springframework.cloud.function.core.FluxToMonoFunction;
+import org.springframework.cloud.function.core.MonoToFluxFunction;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.support.MessageBuilder;
+
+/**
+ * @author Dave Syer
+ */
+public class MessageFunction
+		implements Function<Publisher<Message<?>>, Publisher<Message<?>>> {
+
+	private Function<?, ?> delegate;
+
+	public MessageFunction(Function<?, ?> delegate) {
+		this.delegate = delegate;
+	}
+
+	@Override
+	public Publisher<Message<?>> apply(Publisher<Message<?>> input) {
+		Flux<Message<?>> flux = Flux.from(input);
+		if (this.delegate instanceof FluxFunction) {
+			@SuppressWarnings("unchecked")
+			Function<Object, Object> target = (Function<Object, Object>) ((FluxFunction<?, ?>) this.delegate)
+					.getTarget();
+			return flux.map(
+					value -> MessageBuilder.withPayload(target.apply(value.getPayload()))
+							.copyHeaders(value.getHeaders()).build());
+		}
+		if (this.delegate instanceof MonoToFluxFunction) {
+			@SuppressWarnings("unchecked")
+			Function<Mono<Object>, Flux<Object>> target = ((MonoToFluxFunction<Object, Object>) this.delegate)
+					.getTarget();
+			return flux.next()
+					.flatMapMany(value -> target.apply(Mono.just(value.getPayload()))
+							.map(object -> MessageBuilder.withPayload(object)
+									.copyHeaders(value.getHeaders()).build()));
+		}
+		if (this.delegate instanceof FluxToMonoFunction) {
+			@SuppressWarnings("unchecked")
+			Function<Flux<Object>, Mono<Object>> target = ((FluxToMonoFunction<Object, Object>) this.delegate)
+					.getTarget();
+			AtomicReference<MessageHeaders> headers = new AtomicReference<>();
+			return target.apply(flux.map(messsage -> {
+				headers.set(messsage.getHeaders());
+				return messsage.getPayload();
+			})).map(payload -> MessageBuilder.withPayload(payload)
+					.copyHeaders(headers.get()).build());
+		}
+		if (this.delegate instanceof FluxConsumer) {
+			@SuppressWarnings("unchecked")
+			FluxConsumer<Object> target = ((FluxConsumer<Object>) this.delegate);
+			AtomicReference<MessageHeaders> headers = new AtomicReference<>();
+			Mono<Void> mapped = target.apply(flux.map(messsage -> {
+				headers.set(messsage.getHeaders());
+				return messsage.getPayload();
+			}));
+			return mapped.map(value -> MessageBuilder.createMessage(null, headers.get()));
+		}
+		// TODO: cover the case that delegate is actually Function<Flux,Flux>
+		@SuppressWarnings("unchecked")
+		Function<Object, Object> function = (Function<Object, Object>) this.delegate;
+		return flux.map(
+				value -> MessageBuilder.withPayload(function.apply(value.getPayload()))
+						.copyHeaders(value.getHeaders()).build());
+	}
+
+}

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/MessageSupplier.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2019-2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.cloud.function.context.catalog;
+
+import java.util.function.Supplier;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import org.springframework.cloud.function.core.FluxSupplier;
+import org.springframework.cloud.function.core.MonoSupplier;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.support.MessageBuilder;
+
+/**
+ * @author Dave Syer
+ */
+public class MessageSupplier implements Supplier<Publisher<Message<?>>> {
+
+	private Supplier<?> delegate;
+
+	public MessageSupplier(Supplier<?> delegate) {
+		this.delegate = delegate;
+	}
+
+	@Override
+	public Publisher<Message<?>> get() {
+		if (this.delegate instanceof FluxSupplier) {
+			return ((Flux<?>) this.delegate.get())
+					.map(value -> MessageBuilder.withPayload(value).build());
+		}
+		if (this.delegate instanceof MonoSupplier) {
+			return ((Mono<?>) this.delegate.get())
+					.map(value -> MessageBuilder.withPayload(value).build());
+		}
+		Object product = this.delegate.get();
+		if (product instanceof Publisher) {
+			return Flux.from((Publisher<?>) product)
+					.map(value -> MessageBuilder.withPayload(value).build());
+		}
+		if (product instanceof Iterable) {
+			return Flux.fromIterable((Iterable<?>) product)
+					.map(value -> MessageBuilder.withPayload(value).build());
+		}
+		return Mono.just(MessageBuilder.withPayload(product).build());
+	}
+
+}