How does a flatMap on a double-generic data-structure looks like?

Question

I have the following (simple) data structure:

struct Work<Input, Output> {
    let work: Input -> Output
}

This type represents work which can take an Input and turns in into a desired Output. I am trying to see whether this data structure conforms to some functional concepts like a functor or a monad.

Functor

extension Work {
    func map<MoreOutput>(transform: Output -> MoreOutput) -> Work<Input, MoreOutput> {
        return Work<Input, MoreOutput> {
            return transform(self.work($0))
        }
    }
}

That seems to be correct as far as I am aware. I am able to write a map function which can turn Work<Input, Output> into Work<Input, MoreOutput>

Monad

I have trouble thinking of the definition for a flatMap (or fold) function for Work. The only thing I can come up with is the following:

extension Work {
    func flatMap<MoreOutput>(transform: Work<Output, MoreOutput>) -> Work<Input, MoreOutput> {
        return Work<Input, MoreOutput> { input in
            return transform.work(self.work(input))
        }
    }
}

If you look up the flatMap definition for an Array in swift it looks like this (simplified):

func flatMap(transform: (Element) -> T?) -> [T]

This is a function where its argument is a function which transforms an Element into T and results an Array. I cannot think of a way to abstract this to the Work type.

From another functional book I found a general definition for flatMap as follows (on an object F holding type A):

func flatMap<B>(f: A -> F<B>) -> F<B>

which is a different definition of flatMap than Array seems to implement.

Can someone explain this difference to me? And is it even possible to define a 'correct' flatMap function on Work? Or does Work not satisfy the properties to be a Monad?

** Edit

Thanks phg for so much useful info. I've tried to do the Profunctor definition:

Making Work a Profunctor:

extension Work {
    func diMap<A, B>(fa: A -> Input, fb: Output -> B) -> Work<A, B> {
        return Work<A, B> { arg in
            let input = fa(arg)
            let output = self.work(input)
            return fb(output)
        }
    }
}

Does that look right to you?

Answer 1

This:

func flatMap<B>(f: A -> F<B>) -> F<B>

is what you want flatMap to look like; it's the monad's usual "bind" operation. Specialized for functions over the second argument, you get the so-called Reader monad:

extension Work {
    func flatMap<MoreOutput>(g: Output -> Work<Input, MoreOutput>) -> Work<Input, MoreOutput> {
        // (Reader f) >>= g = Reader $ \x -> runReader (g (f x)) x
        return Work<Input, MoreOutput> {
            g(self.work($0)).work($0)
        }
    }
}

Note: I actually don't speak Swift, this code was just guessing -- hence the included Haskell original. Feel free to edit in a corrected version.

---

Now to the other definition:

func flatMap(transform: (Element) -> T?) -> [T]

I suppose T? means something like "optional T" or "nullable T". This is not what we usually understand as a monadic function, but it is related. Indeed, there has been a question about such "generalized flatMaps". The answer is, that if two monads are compatible, i.e., there exists a monad morphism F<A> -> G<A> preserving monadic structure, it makes sense to define

func wrappedFlatMap<B>(f: A -> F<B>) -> G<B>

which is probably exactly what is happening here for the "option type" and the list type, where the morphism is logically just

Just x ~> [x]
Nothing ~> []