Properly update type context in dataflow pass