Part of [SPARK-2456] Removed some HashMaps from DAGScheduler by storing information in Stage.

This is part of the scheduler cleanup/refactoring effort to make the scheduler code easier to maintain. 

@kayousterhout @markhamstra please take a look ... 
And as a result of the above, you can consolidate this if/else into one thing