Unify treatment of input (stage 1) and output (stage 3)

I was tired of the amount of ifs that changed behaviour
depending on stage 1 or 3. This is a cleaner approach.