TreeGen fix