Fix union type context in type inference

Fixes #332.