Address reviewer comments

I incorporated changes from #358, so that when that PR gets
rebased, it should result in minimal clashes.