Index: src/InitTweak/FixInit.cc
===================================================================
--- src/InitTweak/FixInit.cc	(revision eada3cfc526a731d1ac85003f73c04f68d7ae037)
+++ src/InitTweak/FixInit.cc	(revision c0714bfd6119abb84c6c2aa4881ff5f14758638c)
@@ -70,37 +70,26 @@
 namespace InitTweak {
 	namespace {
-		typedef std::unordered_map< Expression *, TypeSubstitution * > EnvMap;
 		typedef std::unordered_map< int, int > UnqCount;
 
-		class InsertImplicitCalls : public WithTypeSubstitution {
-		public:
+		struct InsertImplicitCalls : public WithTypeSubstitution {
 			/// wrap function application expressions as ImplicitCopyCtorExpr nodes so that it is easy to identify which
 			/// function calls need their parameters to be copy constructed
-			static void insert( std::list< Declaration * > & translationUnit, EnvMap & envMap );
-
-			InsertImplicitCalls( EnvMap & envMap ) : envMap( envMap ) {}
+			static void insert( std::list< Declaration * > & translationUnit );
 
 			Expression * postmutate( ApplicationExpr * appExpr );
-			void premutate( StmtExpr * stmtExpr );
-
-			// collects environments for relevant nodes
-			EnvMap & envMap;
 		};
 
-		class ResolveCopyCtors final : public SymTab::Indexer {
-		public:
+		struct ResolveCopyCtors final : public WithIndexer, public WithShortCircuiting, public WithTypeSubstitution {
 			/// generate temporary ObjectDecls for each argument and return value of each ImplicitCopyCtorExpr,
 			/// generate/resolve copy construction expressions for each, and generate/resolve destructors for both
 			/// arguments and return value temporaries
-			static void resolveImplicitCalls( std::list< Declaration * > & translationUnit, const EnvMap & envMap, UnqCount & unqCount );
-
-			typedef SymTab::Indexer Parent;
-			using Parent::visit;
-
-			ResolveCopyCtors( const EnvMap & envMap, UnqCount & unqCount ) : envMap( envMap ), unqCount( unqCount ) {}
-
-			virtual void visit( ImplicitCopyCtorExpr * impCpCtorExpr ) override;
-			virtual void visit( UniqueExpr * unqExpr ) override;
-			virtual void visit( StmtExpr * stmtExpr ) override;
+			static void resolveImplicitCalls( std::list< Declaration * > & translationUnit, UnqCount & unqCount );
+
+			ResolveCopyCtors( UnqCount & unqCount ) : unqCount( unqCount ) {}
+
+			void postvisit( ImplicitCopyCtorExpr * impCpCtorExpr );
+			void postvisit( StmtExpr * stmtExpr );
+			void previsit( UniqueExpr * unqExpr );
+			void postvisit( UniqueExpr * unqExpr );
 
 			/// create and resolve ctor/dtor expression: fname(var, [cpArg])
@@ -111,7 +100,6 @@
 			void destructRet( ObjectDecl * ret, ImplicitCopyCtorExpr * impCpCtorExpr );
 
-			TypeSubstitution * env;
-			const EnvMap & envMap;
 			UnqCount & unqCount; // count the number of times each unique expr ID appears
+			std::unordered_set< int > vars;
 		};
 
@@ -292,9 +280,8 @@
 		InitTweak::fixGlobalInit( translationUnit, filename, inLibrary );
 
-		EnvMap envMap;
 		UnqCount unqCount;
 
-		InsertImplicitCalls::insert( translationUnit, envMap );
-		ResolveCopyCtors::resolveImplicitCalls( translationUnit, envMap, unqCount );
+		InsertImplicitCalls::insert( translationUnit );
+		ResolveCopyCtors::resolveImplicitCalls( translationUnit, unqCount );
 		InsertDtors::insert( translationUnit );
 		FixInit::fixInitializers( translationUnit );
@@ -315,11 +302,11 @@
 
 	namespace {
-		void InsertImplicitCalls::insert( std::list< Declaration * > & translationUnit, EnvMap & envMap ) {
-			PassVisitor<InsertImplicitCalls> inserter( envMap );
+		void InsertImplicitCalls::insert( std::list< Declaration * > & translationUnit ) {
+			PassVisitor<InsertImplicitCalls> inserter;
 			mutateAll( translationUnit, inserter );
 		}
 
-		void ResolveCopyCtors::resolveImplicitCalls( std::list< Declaration * > & translationUnit, const EnvMap & envMap, UnqCount & unqCount ) {
-			ResolveCopyCtors resolver( envMap, unqCount );
+		void ResolveCopyCtors::resolveImplicitCalls( std::list< Declaration * > & translationUnit, UnqCount & unqCount ) {
+			PassVisitor<ResolveCopyCtors> resolver( unqCount );
 			acceptAll( translationUnit, resolver );
 		}
@@ -395,20 +382,11 @@
 			// wrap each function call so that it is easy to identify nodes that have to be copy constructed
 			ImplicitCopyCtorExpr * expr = new ImplicitCopyCtorExpr( appExpr );
-			// save the type substitution into the envMap so that it is easy to find.
+			// Move the type substitution to the new top-level, if it is attached to the appExpr.
 			// Ensure it is not deleted with the ImplicitCopyCtorExpr by removing it before deletion.
 			// The substitution is needed to obtain the type of temporary variables so that copy constructor
-			// calls can be resolved. Normally this is what PolyMutator is for, but the pass that resolves
-			// copy constructor calls must be an Indexer. We could alternatively make a PolyIndexer which
-			// saves the environment, or compute the types of temporaries here, but it's much simpler to
-			// save the environment here, and more cohesive to compute temporary variables and resolve copy
-			// constructor calls together.
+			// calls can be resolved.
 			assert( env );
-			envMap[expr] = env;
+			std::swap( expr->env, appExpr->env );
 			return expr;
-		}
-
-		void InsertImplicitCalls::premutate( StmtExpr * stmtExpr ) {
-			assert( env );
-			envMap[stmtExpr] = env;
 		}
 
@@ -428,5 +406,5 @@
 			// (VariableExpr and already resolved expression)
 			CP_CTOR_PRINT( std::cerr << "ResolvingCtorDtor " << untyped << std::endl; )
-			Expression * resolved = ResolvExpr::findVoidExpression( untyped, *this );
+			Expression * resolved = ResolvExpr::findVoidExpression( untyped, indexer );
 			assert( resolved );
 			if ( resolved->get_env() ) {
@@ -477,9 +455,6 @@
 		}
 
-		void ResolveCopyCtors::visit( ImplicitCopyCtorExpr *impCpCtorExpr ) {
+		void ResolveCopyCtors::postvisit( ImplicitCopyCtorExpr *impCpCtorExpr ) {
 			CP_CTOR_PRINT( std::cerr << "ResolveCopyCtors: " << impCpCtorExpr << std::endl; )
-			Parent::visit( impCpCtorExpr );
-			env = envMap.at(impCpCtorExpr);
-			assert( env );
 
 			ApplicationExpr * appExpr = impCpCtorExpr->get_callExpr();
@@ -510,7 +485,6 @@
 		}
 
-		void ResolveCopyCtors::visit( StmtExpr * stmtExpr ) {
-			Parent::visit( stmtExpr );
-			env = envMap.at(stmtExpr);
+		void ResolveCopyCtors::postvisit( StmtExpr * stmtExpr ) {
+			assert( env );
 			assert( stmtExpr->get_result() );
 			Type * result = stmtExpr->get_result();
@@ -534,16 +508,20 @@
 				stmtExpr->get_dtors().push_front( makeCtorDtor( "^?{}", ret ) );
 			} // if
-
-		}
-
-		void ResolveCopyCtors::visit( UniqueExpr * unqExpr ) {
-			static std::unordered_set< int > vars;
+		}
+
+		void ResolveCopyCtors::previsit( UniqueExpr * unqExpr ) {
 			unqCount[ unqExpr->get_id() ]++;  // count the number of unique expressions for each ID
 			if ( vars.count( unqExpr->get_id() ) ) {
 				// xxx - hack to prevent double-handling of unique exprs, otherwise too many temporary variables and destructors are generated
+				visit_children = false;
+			}
+		}
+
+		void ResolveCopyCtors::postvisit( UniqueExpr * unqExpr ) {
+			if ( vars.count( unqExpr->get_id() ) ) {
+				// xxx - hack to prevent double-handling of unique exprs, otherwise too many temporary variables and destructors are generated
 				return;
 			}
 
-			Parent::visit( unqExpr );
 			// it should never be necessary to wrap a void-returning expression in a UniqueExpr - if this assumption changes, this needs to be rethought
 			assert( unqExpr->get_result() );
@@ -582,5 +560,5 @@
 
 			// xxx - update to work with multiple return values
-			ObjectDecl * returnDecl = returnDecls.empty() ? NULL : returnDecls.front();
+			ObjectDecl * returnDecl = returnDecls.empty() ? nullptr : returnDecls.front();
 			Expression * callExpr = impCpCtorExpr->get_callExpr();
 
@@ -591,6 +569,7 @@
 			tempDecls.clear();
 			returnDecls.clear();
-			impCpCtorExpr->set_callExpr( NULL );
-			impCpCtorExpr->set_env( NULL );
+			impCpCtorExpr->set_callExpr( nullptr );
+			std::swap( impCpCtorExpr->env, callExpr->env );
+			assert( impCpCtorExpr->env == nullptr );
 			delete impCpCtorExpr;
 
