Flatten shapes when fetching with Model.shape

gems/smithy/lib/smithy/model.rb

@@ -36,16 +36,99 @@ module Model
       'smithy.api#Unit' => { 'type' => 'structure', 'members' => {}, 'traits' => { 'smithy.api#unitType' => {} } }
     }.freeze
 
-    # @param [Hash] model Model
-    # @param [String] target Target shape
-    # @return [Hash] The shape
-    def self.shape(model, target)
-      if model['shapes'].key?(target)
-        model['shapes'][target]
-      elsif PRELUDE_SHAPES.key?(target)
-        PRELUDE_SHAPES[target]
-      else
-        raise ArgumentError, "Shape not found: #{target}"
+    class << self
+      # @param [Hash] model Model
+      # @param [String] target Target shape
+      # @return [Hash] The shape
+      def shape(model, target)
+        if model['shapes'].key?(target)
+          model['shapes'][target]
+        elsif PRELUDE_SHAPES.key?(target)
+          PRELUDE_SHAPES[target]
+        else
+          raise ArgumentError, "Shape not found: #{target}"
+        end
+      end
+
+      # resolve mixins and "apply" types.
+      def preprocess(model)
+        # TODO: Deep dup of model
+        model['shapes'].each do |id, shape|
+          flatten_shape(id, model, shape)
+        end
+        model
+      end
+
+      private
+
+      def flatten_shape(id, model, shape)
+        mixin_members, mixin_traits = resolve_mixins(model, shape)
+
+        shape['members'] = mixin_members.merge(shape.fetch('members', {})) unless mixin_members.empty?
+        shape['traits'] = merge2(mixin_traits, shape.fetch('traits', {})) unless mixin_traits.empty?
+
+        shape.delete('mixins') # remove mixins now that we've resolved them all
+
+        # after mixins are done, then check for any "apply" for any members
+        flatten_apply_traits(id, model, shape)
+
+        shape
+      end
+
+      def flatten_apply_traits(id, model, shape)
+        shape.fetch('members', {}).each_key do |name|
+          # check for an apply
+          member_id = "#{id}$#{name}"
+          next unless model['shapes'][member_id] && model['shapes'][member_id]['type'] == 'apply'
+
+          flatten_apply_trait(model, shape, name, member_id)
+        end
+      end
+
+      def flatten_apply_trait(model, shape, name, member_id)
+        shape['members'][name]['traits'] =
+          shape['members'][name].fetch('traits', {}).merge(model['shapes'][member_id]['traits'])
+
+        model['shapes'].delete(member_id)
+      end
+
+      def resolve_mixins(model, shape)
+        mixin_traits = {}
+        mixin_members = {}
+        shape.fetch('mixins', []).each do |mixin|
+          mixin_shape = flatten_shape(mixin['target'], model, model['shapes'][mixin['target']])
+          mixin_members = mixin_members.merge(mixin_shape['members']) if mixin_shape['members']
+
+          next unless mixin_shape['traits']
+
+          mixin_traits = resolve_mixin_traits(mixin_shape, mixin_traits)
+        end
+
+        [mixin_members, mixin_traits]
+      end
+
+      def resolve_mixin_traits(mixin_shape, mixin_traits)
+        skip_traits = mixin_shape['traits'].fetch('smithy.api#mixin', {}).fetch('localTraits', [])
+        skip_traits << 'smithy.api#mixin'
+        merge2(
+          mixin_traits,
+          mixin_shape['traits'].except(*skip_traits)
+        )
+      end
+
+      # A variant of Hash.merge that places keys from hash2 first rather than last
+      # Order of precedence is the same as Hash.merge.
+      # values from hash2 are used over values from hash1.
+      # Keys in hash2 are first in the resulting hash.
+      def merge2(hash1, hash2)
+        result = {}
+        hash2.each do |key, value|
+          result[key] = value
+        end
+        hash1.each do |key, value|
+          result[key] = value unless result.key?(key)
+        end
+        result
       end
     end
   end

gems/smithy/lib/smithy/plan.rb

@@ -16,17 +16,11 @@ class Plan
     #  If not provided, source code will be generated to STDOUT.
     # @option options [Boolean] :quiet Whether to suppress output.
     def initialize(model, type, options = {})
-      @model = model
+      @model = Model.preprocess(model)
       @type = type
       @service = find_service(model['shapes'])
 
-      @name = options.fetch(:name, default_name(service))
-      @module_name = options.fetch(:module_name, @name)
-      @gem_name = options.fetch(:gem_name, default_gem_name(@module_name, @type))
-      @gem_version = options.fetch(:gem_version)
-
-      @destination_root = options.fetch(:destination_root, nil)
-      @quiet = options.fetch(:quiet, false)
+      initialize_options(options)
 
       Welds.load!(self)
       @welds = Welds.for(@service)
@@ -64,6 +58,16 @@ def initialize(model, type, options = {})
 
     private
 
+    def initialize_options(options)
+      @name = options.fetch(:name, default_name(service))
+      @module_name = options.fetch(:module_name, @name)
+      @gem_name = options.fetch(:gem_name, default_gem_name(@module_name, @type))
+      @gem_version = options.fetch(:gem_version)
+
+      @destination_root = options.fetch(:destination_root, nil)
+      @quiet = options.fetch(:quiet, false)
+    end
+
     def find_service(shapes)
       service = shapes.select { |_, shape| shape['type'] == 'service' }
       raise 'Multiple service shapes found' if service.size > 1