Cleanup/fix some logic in lib.