Add type sigs and minor refactoring