Also refactor to simplify data structure Test: atest Bug: 149386091 Change-Id: Iccc8ee440bdd02a74390713c3dfeffa49d1ba142