Restructure and add in Traits