Minor cleanup of original patch to change less code.