Clean up treegen code

Created before #308, but I never got around to merging it apparently.
Both pulls include changes worthy of merging.