Better types, more succinct code

Change the type to reflect what is going on in the code.